fix: updated based on review

Author: abhishekmadan30

docs/concepts/transforms.rst

@@ -110,7 +110,7 @@ about the state of the tasks at given points. Here is an example:
    from taskgraph.util.schema import Schema
 
    class MySchema(Schema):
-       foo: str                   # Required field
+       foo: str  # Required field
        bar: Optional[bool] = None # Optional field
 
    transforms = TransformSequence()

src/taskgraph/config.py

@@ -23,7 +23,7 @@
 ]
 
 
-class WorkerAlias(Schema):
+class WorkerAliasSchema(Schema):
     """Worker alias configuration."""
 
     provisioner: optionally_keyed_by("level", str)  # type: ignore
@@ -32,10 +32,10 @@ class WorkerAlias(Schema):
     worker_type: optionally_keyed_by("level", str)  # type: ignore
 
 
-class Workers(Schema, rename=None):
+class WorkersSchema(Schema, rename=None):
     """Workers configuration."""
 
-    aliases: Dict[str, WorkerAlias]
+    aliases: Dict[str, WorkerAliasSchema]
 
 
 class Repository(Schema, forbid_unknown_fields=False):
@@ -58,7 +58,7 @@ class RunConfig(Schema):
     use_caches: Optional[Union[bool, List[str]]] = None  # Maps from "use-caches"
 
 
-class TaskGraphConfig(Schema):
+class TaskGraphSchema(Schema):
     """Taskgraph specific configuration."""
 
     # Required fields first
@@ -82,8 +82,8 @@ class GraphConfigSchema(Schema, forbid_unknown_fields=False):
     # Required fields first
     trust_domain: str  # Maps from "trust-domain"
     task_priority: optionally_keyed_by("project", "level", TaskPriority)  # type: ignore
-    workers: Workers
-    taskgraph: TaskGraphConfig
+    workers: WorkersSchema
+    taskgraph: TaskGraphSchema
 
     # Optional fields
     docker_image_kind: Optional[str] = None  # Maps from "docker-image-kind"
@@ -158,7 +158,6 @@ def kinds_dir(self):
 
 def validate_graph_config(config):
     """Validate graph configuration using msgspec."""
-    # With rename="kebab", msgspec handles the conversion automatically
     validate_schema(GraphConfigSchema, config, "Invalid graph configuration:")
 
 

src/taskgraph/parameters.py

@@ -29,7 +29,7 @@ class ParameterMismatch(Exception):
     """Raised when a parameters.yml has extra or missing parameters."""
 
 
-class CodeReviewConfig(Schema):
+class CodeReviewSchema(Schema):
     """Code review configuration."""
 
     # Required field
@@ -73,10 +73,10 @@ class BaseSchema(Schema):
     tasks_for: str
 
     # Optional fields
-    next_version: Optional[str]
-    optimize_strategies: Optional[str]
-    version: Optional[str]
-    code_review: Optional[CodeReviewConfig] = None
+    next_version: Optional[str] = None
+    optimize_strategies: Optional[str] = None
+    version: Optional[str] = None
+    code_review: Optional[CodeReviewSchema] = None
 
 
 def get_contents(path):
@@ -243,30 +243,74 @@ def check(self):
 
             if self.strict:
                 # Strict mode: validate against schema and check for extra fields
-                # Get all valid field names from the base schema
+                # Get all valid field names from the base schema and extensions
                 schema_fields = {
                     f.encode_name for f in msgspec.structs.fields(BaseSchema)
                 }
 
+                # Add fields from extension schemas
+                for ext_schema in _schema_extensions:
+                    if isinstance(ext_schema, type) and issubclass(
+                        ext_schema, msgspec.Struct
+                    ):
+                        schema_fields.update(
+                            {f.encode_name for f in msgspec.structs.fields(ext_schema)}
+                        )
+
                 # Check for extra fields
                 extra_fields = set(kebab_params.keys()) - schema_fields
                 if extra_fields:
                     raise ParameterMismatch(
                         f"Invalid parameters: Extra fields not allowed: {extra_fields}"
                     )
 
-                # Validate all parameters against the schema
-                msgspec.convert(kebab_params, BaseSchema)
+                # Validate base schema fields only (filter out extension fields)
+                base_fields = {
+                    f.encode_name for f in msgspec.structs.fields(BaseSchema)
+                }
+                base_params = {
+                    k: v for k, v in kebab_params.items() if k in base_fields
+                }
+                msgspec.convert(base_params, BaseSchema)
+
+                # Also validate against extension schemas
+                for ext_schema in _schema_extensions:
+                    if isinstance(ext_schema, type) and issubclass(
+                        ext_schema, msgspec.Struct
+                    ):
+                        # Only validate fields that belong to this extension
+                        ext_fields = {
+                            f.encode_name for f in msgspec.structs.fields(ext_schema)
+                        }
+                        ext_params = {
+                            k: v for k, v in kebab_params.items() if k in ext_fields
+                        }
+                        if ext_params:
+                            msgspec.convert(ext_params, ext_schema)
             else:
-                # Non-strict mode: only validate fields that exist in the schema
-                # Filter to only include fields defined in the schema
+                # Non-strict mode: only validate fields that exist in the schemas
+                # Filter to only include fields defined in the base schema
                 schema_fields = {
                     f.encode_name for f in msgspec.structs.fields(BaseSchema)
                 }
                 filtered_params = {
                     k: v for k, v in kebab_params.items() if k in schema_fields
                 }
                 msgspec.convert(filtered_params, BaseSchema)
+
+                # Also validate extension schemas in non-strict mode
+                for ext_schema in _schema_extensions:
+                    if isinstance(ext_schema, type) and issubclass(
+                        ext_schema, msgspec.Struct
+                    ):
+                        ext_fields = {
+                            f.encode_name for f in msgspec.structs.fields(ext_schema)
+                        }
+                        ext_params = {
+                            k: v for k, v in kebab_params.items() if k in ext_fields
+                        }
+                        if ext_params:
+                            msgspec.convert(ext_params, ext_schema)
         except (msgspec.ValidationError, msgspec.DecodeError) as e:
             raise ParameterMismatch(f"Invalid parameters: {e}")
 

src/taskgraph/transforms/chunking.py

@@ -9,7 +9,7 @@
 from taskgraph.util.templates import substitute
 
 
-class ChunkConfig(Schema):
+class ChunkSchema(Schema):
     """
     `chunk` can be used to split one task into `total-chunks`
     tasks, substituting `this_chunk` and `total_chunks` into any
@@ -24,13 +24,13 @@ class ChunkConfig(Schema):
 
 
 #: Schema for chunking transforms
-class ChunkSchema(Schema, forbid_unknown_fields=False):
+class ChunksSchema(Schema, forbid_unknown_fields=False):
     # Optional, so it can be used for a subset of tasks in a kind
-    chunk: Optional[ChunkConfig] = None
+    chunk: Optional[ChunkSchema] = None
 
 
 transforms = TransformSequence()
-transforms.add_validate(ChunkSchema)
+transforms.add_validate(ChunksSchema)
 
 
 @transforms.add

src/taskgraph/transforms/fetch.py

@@ -9,7 +9,7 @@
 import os
 import re
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, Optional, Union
 
 import msgspec
 
@@ -24,14 +24,6 @@
 CACHE_TYPE = "content.v1"
 
 
-#: Schema for fetch transforms
-class FetchConfig(Schema, rename=None, omit_defaults=False):
-    """Configuration for a fetch task type."""
-
-    type: str
-    # Additional fields handled dynamically by fetch builders
-
-
 class FetchSchema(Schema):
     # Required fields
     # Name of the task.
@@ -70,7 +62,7 @@ def __post_init__(self):
 
 @dataclass(frozen=True)
 class FetchBuilder:
-    schema: Any  # Either msgspec.Struct type or validation function
+    schema: Union[Schema, Callable]
     builder: Callable
 
 
@@ -194,7 +186,7 @@ def make_task(config, tasks):
         yield task_desc
 
 
-class GPGSignatureConfig(Schema):
+class GPGSignatureSchema(Schema):
     """GPG signature verification configuration."""
 
     # URL where GPG signature document can be obtained. Can contain the
@@ -204,7 +196,7 @@ class GPGSignatureConfig(Schema):
     key_path: str
 
 
-class StaticUrlFetchConfig(Schema, rename="kebab"):
+class StaticUrlFetchSchema(Schema):
     """Configuration for static-url fetch type."""
 
     type: str
@@ -215,7 +207,7 @@ class StaticUrlFetchConfig(Schema, rename="kebab"):
     # Size of the downloaded entity, in bytes.
     size: int
     # GPG signature verification.
-    gpg_signature: Optional[GPGSignatureConfig] = None
+    gpg_signature: Optional[GPGSignatureSchema] = None
     # The name to give to the generated artifact. Defaults to the file
     # portion of the URL. Using a different extension converts the
     # archive to the given type. Only conversion to .tar.zst is supported.
@@ -233,7 +225,7 @@ class StaticUrlFetchConfig(Schema, rename="kebab"):
     # it is important to update the digest data used to compute cache hits.
 
 
-@fetch_builder("static-url", StaticUrlFetchConfig)
+@fetch_builder("static-url", StaticUrlFetchSchema)
 def create_fetch_url_task(config, name, fetch):
     artifact_name = fetch.get("artifact-name")
     if not artifact_name:
@@ -296,7 +288,7 @@ def create_fetch_url_task(config, name, fetch):
     }
 
 
-class GitFetchConfig(Schema):
+class GitFetchSchema(Schema):
     """Configuration for git fetch type."""
 
     type: str
@@ -312,7 +304,7 @@ class GitFetchConfig(Schema):
     ssh_key: Optional[str] = None
 
 
-@fetch_builder("git", GitFetchConfig)
+@fetch_builder("git", GitFetchSchema)
 def create_git_fetch_task(config, name, fetch):
     path_prefix = fetch.get("path-prefix")
     if not path_prefix:

src/taskgraph/transforms/from_deps.py

@@ -13,25 +13,19 @@
 
 from copy import deepcopy
 from textwrap import dedent
-from typing import Any, Dict, List, Optional, Union
-
-import msgspec
+from typing import Any, Dict, List, Literal, Optional, Union
 
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.util.attributes import attrmatch
 from taskgraph.util.dependencies import GROUP_BY_MAP, get_dependencies
 from taskgraph.util.schema import Schema, validate_schema
 from taskgraph.util.set_name import SET_NAME_MAP
 
+SetNameType = Literal["strip-kind", "retain-kind"]
+GroupByType = Literal["single", "all", "attribute"]
 
-class FetchEntry(Schema, rename=None):
-    """A fetch entry for an artifact."""
-
-    artifact: str
-    dest: Optional[str] = None
 
-
-class FromDepsConfig(Schema):
+class FromDepsChildSchema(Schema):
     # Optional fields
     # Limit dependencies to specified kinds (defaults to all kinds in
     # `kind-dependencies`).
@@ -40,16 +34,17 @@ class FromDepsConfig(Schema):
     # dependency of this kind will be used to derive the label
     # and copy attributes (if `copy-attributes` is True).
     kinds: Optional[List[str]] = None
-    # UPDATE ME AND DOCS
-    set_name: Optional[Union[str, bool, Dict[str, Any]]] = None
+    # Set the task name using the specified function. Can be False to
+    # disable name setting, or a string/dict specifying the function to use.
+    set_name: Optional[Union[SetNameType, bool, Dict[SetNameType, Any]]] = None
     # Limit dependencies to tasks whose attributes match
     # using :func:`~taskgraph.util.attributes.attrmatch`.
     with_attributes: Optional[Dict[str, Union[List[Any], str]]] = None
     # Group cross-kind dependencies using the given group-by
     # function. One task will be created for each group. If not
     # specified, the 'single' function will be used which creates
     # a new task for each individual dependency.
-    group_by: Optional[Union[str, Dict[str, Any]]] = None
+    group_by: Optional[Union[GroupByType, Dict[GroupByType, Any]]] = None
     # If True, copy attributes from the dependency matching the
     # first kind in the `kinds` list (whether specified explicitly
     # or taken from `kind-dependencies`).
@@ -64,35 +59,12 @@ class FromDepsConfig(Schema):
     # `fetches` entry.
     fetches: Optional[Dict[str, List[Union[str, Dict[str, str]]]]] = None
 
-    def __post_init__(self):
-        # Validate set_name
-        if self.set_name is not None and self.set_name is not False:
-            if isinstance(self.set_name, str) and self.set_name not in SET_NAME_MAP:
-                raise msgspec.ValidationError(f"Invalid set-name: {self.set_name}")
-            elif isinstance(self.set_name, dict):
-                keys = list(self.set_name.keys())
-                if len(keys) != 1 or keys[0] not in SET_NAME_MAP:
-                    raise msgspec.ValidationError(
-                        f"Invalid set-name dict: {self.set_name}"
-                    )
-
-        # Validate group_by
-        if self.group_by is not None:
-            if isinstance(self.group_by, str) and self.group_by not in GROUP_BY_MAP:
-                raise msgspec.ValidationError(f"Invalid group-by: {self.group_by}")
-            elif isinstance(self.group_by, dict):
-                keys = list(self.group_by.keys())
-                if len(keys) != 1 or keys[0] not in GROUP_BY_MAP:
-                    raise msgspec.ValidationError(
-                        f"Invalid group-by dict: {self.group_by}"
-                    )
-
-
-#: Schema for from_deps transforms
+
+# Schema for from_deps transforms
 class FromDepsSchema(Schema, forbid_unknown_fields=False):
     """Schema for from_deps transforms."""
 
-    from_deps: FromDepsConfig
+    from_deps: FromDepsChildSchema
 
 
 transforms = TransformSequence()

src/taskgraph/transforms/matrix.py

@@ -15,7 +15,7 @@
 from taskgraph.util.templates import substitute_task_fields
 
 
-class MatrixConfig(Schema, forbid_unknown_fields=False):
+class MatrixChildSchema(Schema, forbid_unknown_fields=False):
     """
     Matrix configuration for generating multiple tasks.
     """
@@ -47,7 +47,7 @@ class MatrixSchema(Schema, forbid_unknown_fields=False):
     """
 
     name: str
-    matrix: Optional[MatrixConfig] = None
+    matrix: Optional[MatrixChildSchema] = None
 
 
 transforms = TransformSequence()