fix: changed how extra fields work

Author: abhishekmadan30

src/taskgraph/config.py

@@ -8,9 +8,7 @@
 import sys
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, List, Literal, Optional, Union
-
-import msgspec
+from typing import Dict, List, Literal, Optional, Union
 
 from .util.python_path import find_object
 from .util.schema import Schema, optionally_keyed_by, validate_schema
@@ -55,17 +53,18 @@ class Workers(Schema, rename=None):
     aliases: Dict[str, WorkerAlias]
 
 
-class Repository(Schema):
-    """Repository configuration."""
+class Repository(Schema, forbid_unknown_fields=False):
+    """Repository configuration.
+
+    This schema allows extra fields for repository-specific configuration.
+    """
 
     # Required fields first
     name: str
 
     # Optional fields
     project_regex: Optional[str] = None  # Maps from "project-regex"
     ssh_secret_name: Optional[str] = None  # Maps from "ssh-secret-name"
-    # Allow extra fields for flexibility
-    __extras__: Dict[str, Any] = msgspec.field(default_factory=dict)
 
 
 class RunConfig(Schema):
@@ -89,8 +88,11 @@ class TaskGraphConfig(Schema):
     run: Optional[RunConfig] = None
 
 
-class GraphConfigSchema(Schema):
-    """Main graph configuration schema."""
+class GraphConfigSchema(Schema, forbid_unknown_fields=False):
+    """Main graph configuration schema.
+
+    This schema allows extra fields for flexibility in graph configuration.
+    """
 
     # Required fields first
     trust_domain: str  # Maps from "trust-domain"
@@ -106,8 +108,6 @@ class GraphConfigSchema(Schema):
         None  # Maps from "task-deadline-after", can be keyed-by project
     )
     task_expires_after: Optional[str] = None  # Maps from "task-expires-after"
-    # Allow extra fields for flexibility
-    __extras__: Dict[str, Any] = msgspec.field(default_factory=dict)
 
     def __post_init__(self):
         """Validate keyed-by fields."""

src/taskgraph/transforms/chunking.py

@@ -2,9 +2,7 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 import copy
-from typing import Any, Dict, List, Optional
-
-import msgspec
+from typing import List, Optional
 
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.util.schema import Schema
@@ -26,16 +24,13 @@ class ChunkConfig(Schema):
 
 
 #: Schema for chunking transforms
-class ChunkSchema(Schema):
+class ChunkSchema(Schema, forbid_unknown_fields=False):
     # Optional, so it can be used for a subset of tasks in a kind
     chunk: Optional[ChunkConfig] = None
-    __extras__: Dict[str, Any] = msgspec.field(default_factory=dict)
-
 
-CHUNK_SCHEMA = ChunkSchema
 
 transforms = TransformSequence()
-transforms.add_validate(CHUNK_SCHEMA)
+transforms.add_validate(ChunkSchema)
 
 
 @transforms.add

src/taskgraph/transforms/from_deps.py

@@ -89,12 +89,10 @@ def __post_init__(self):
 
 
 #: Schema for from_deps transforms
-class FromDepsSchema(Schema):
+class FromDepsSchema(Schema, forbid_unknown_fields=False):
     """Schema for from_deps transforms."""
 
     from_deps: FromDepsConfig
-    # Allow extra fields
-    _extra: Optional[Dict[str, Any]] = msgspec.field(default=None, name="")
 
 
 # Backward compatibility

src/taskgraph/transforms/matrix.py

@@ -8,16 +8,14 @@
 """
 
 from copy import deepcopy
-from typing import Any, Dict, List, Optional
-
-import msgspec
+from typing import Dict, List, Optional
 
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.util.schema import Schema
 from taskgraph.util.templates import substitute_task_fields
 
 
-class MatrixConfig(Schema):
+class MatrixConfig(Schema, forbid_unknown_fields=False):
     """
     Matrix configuration for generating multiple tasks.
     """
@@ -39,15 +37,17 @@ class MatrixConfig(Schema):
     # If not specified, all fields in the task definition will be
     # substituted.
     substitution_fields: Optional[List[str]] = None
-    # Allow extra fields for matrix dimensions
-    __extras__: Dict[str, List[str]] = msgspec.field(default_factory=dict)
 
 
 #: Schema for matrix transforms
-class MatrixSchema(Schema):
+class MatrixSchema(Schema, forbid_unknown_fields=False):
+    """Schema for matrix transforms.
+
+    This schema allows extra fields to be passed through to the task.
+    """
+
     name: str
     matrix: Optional[MatrixConfig] = None
-    __extras__: Dict[str, Any] = msgspec.field(default_factory=dict)
 
 
 MATRIX_SCHEMA = MatrixSchema

src/taskgraph/transforms/notify.py

@@ -128,7 +128,7 @@ class LegacyNotificationsConfig(Schema, rename="kebab"):
 
 
 #: Schema for notify transforms
-class NotifySchema(Schema, tag_field="notify_type"):
+class NotifySchema(Schema, tag_field="notify_type", forbid_unknown_fields=False):
     """Schema for notify transforms.
 
     Note: This schema allows either 'notify' or 'notifications' field,
@@ -137,8 +137,6 @@ class NotifySchema(Schema, tag_field="notify_type"):
 
     notify: Optional[NotifyConfig] = None
     notifications: Optional[LegacyNotificationsConfig] = None
-    # Allow extra fields
-    _extra: Optional[Dict[str, Any]] = msgspec.field(default=None, name="")
 
     def __post_init__(self):
         # Ensure only one of notify or notifications is present

src/taskgraph/transforms/run/__init__.py

@@ -46,13 +46,14 @@ class WhenConfig(Schema):
 
 
 # Run configuration using msgspec
-class RunConfig(Schema, rename=None):
-    """Configuration for how to run a task."""
+class RunConfig(Schema, rename=None, forbid_unknown_fields=False):
+    """Configuration for how to run a task.
+
+    This schema allows extra fields for run implementation-specific configuration.
+    """
 
     using: str
     workdir: TOptional[str] = None
-    # Allow any extra fields for run implementation-specific config
-    __extras__: Dict[str, Any] = msgspec.field(default_factory=dict)
 
 
 # Run description schema using msgspec

src/taskgraph/transforms/run/toolchain.py

@@ -7,8 +7,6 @@
 
 from typing import Any, Dict, List, Literal, Optional, Union
 
-import msgspec
-
 import taskgraph
 from taskgraph.transforms.run import configure_taskdesc_for_run, run_task_using
 from taskgraph.transforms.run.common import (
@@ -25,7 +23,7 @@
 
 
 #: Schema for run.using toolchain
-class ToolchainRunSchema(Schema):
+class ToolchainRunSchema(Schema, forbid_unknown_fields=False):
     # Required fields first
 
     # Specifies the run type. Must be "toolchain-script".
@@ -63,12 +61,6 @@ class ToolchainRunSchema(Schema):
     # toolchain.
     toolchain_env: Optional[Dict[str, Any]] = None
 
-    # Allow extra fields
-    _extra: Optional[Dict[str, Any]] = msgspec.field(default=None, name="")
-
-
-toolchain_run_schema = ToolchainRunSchema
-
 
 def get_digest_data(config, run, taskdesc):
     files = list(run.pop("resources", []))

src/taskgraph/transforms/task.py

@@ -70,12 +70,13 @@ class TaskDescriptionIndex(Schema, rename="kebab"):
     rank: Union[Literal["by-tier", "build_date"], int] = "by-tier"
 
 
-class TaskDescriptionWorker(Schema, rename=None):
-    """Worker configuration for a task."""
+class TaskDescriptionWorker(Schema, rename=None, forbid_unknown_fields=False):
+    """Worker configuration for a task.
+
+    This schema allows extra fields for worker-specific configuration.
+    """
 
     implementation: str
-    # Allow any extra fields for worker-specific configuration
-    __extras__: Dict[str, TAny] = msgspec.field(default_factory=dict)
 
 
 class TaskDescriptionSchema(Schema):

src/taskgraph/transforms/task_context.py

@@ -1,7 +1,5 @@
 from typing import Any, Dict, List, Optional, Union
 
-import msgspec
-
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.util.schema import Schema
 from taskgraph.util.templates import deep_get, substitute_task_fields
@@ -53,13 +51,17 @@ class TaskContextConfig(Schema):
 
 
 #: Schema for the task_context transforms
-class TaskContextSchema(Schema):
+class TaskContextSchema(Schema, forbid_unknown_fields=False):
+    """Schema for task context transforms.
+
+    This schema allows extra fields to be passed through to the task.
+    """
+
     # Required field first
     task_context: TaskContextConfig
 
     # Optional fields
     name: Optional[str] = None
-    __extras__: Dict[str, Any] = msgspec.field(default_factory=dict)
 
 
 SCHEMA = TaskContextSchema

src/taskgraph/util/schema.py

@@ -190,7 +190,13 @@ def resolve_keyed_by(
 ]
 
 
-class Schema(msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"):
+class Schema(
+    msgspec.Struct,
+    kw_only=True,
+    omit_defaults=True,
+    rename="kebab",
+    forbid_unknown_fields=True,
+):
     """
     Base schema class that extends msgspec.Struct.
 
@@ -202,6 +208,13 @@ class MySchema(Schema):
 
     Instead of wrapping msgspec.Struct types.
     Most schemas use kebab-case renaming by default.
+
+    By default, forbid_unknown_fields is True, meaning extra fields
+    will cause validation errors. Child classes can override this by
+    setting forbid_unknown_fields=False in their class definition:
+
+        class MySchema(Schema, forbid_unknown_fields=False):
+            foo: str
     """
 
     @classmethod