feat: added support for optional keying

Author: abhishekmadan30

src/taskgraph/config.py

@@ -13,7 +13,7 @@
 import msgspec
 
 from .util.python_path import find_object
-from .util.schema import Schema, validate_schema
+from .util.schema import Schema, optionally_keyed_by, validate_schema
 from .util.vcs import get_repository
 from .util.yaml import load_yaml
 
@@ -29,10 +29,24 @@
 class WorkerAlias(Schema):
     """Worker alias configuration."""
 
-    provisioner: Union[str, dict]
+    provisioner: Union[str, dict]  # Can be keyed-by level
     implementation: str
     os: str
-    worker_type: Union[str, dict]  # Can be keyed-by, maps from "worker-type"
+    worker_type: Union[str, dict]  # Can be keyed-by level, maps from "worker-type"
+
+    def __post_init__(self):
+        """Validate keyed-by fields."""
+        # Validate provisioner can be keyed-by level
+        if isinstance(self.provisioner, dict):
+            validator = optionally_keyed_by("level", str)
+            # Just validate - it will raise an error if invalid
+            validator(self.provisioner)
+
+        # Validate worker_type can be keyed-by level
+        if isinstance(self.worker_type, dict):
+            validator = optionally_keyed_by("level", str)
+            # Just validate - it will raise an error if invalid
+            validator(self.worker_type)
 
 
 class Workers(Schema, rename=None):
@@ -82,19 +96,48 @@ class GraphConfigSchema(Schema):
     trust_domain: str  # Maps from "trust-domain"
     task_priority: Union[
         TaskPriority, dict
-    ]  # Maps from "task-priority", can be keyed-by
+    ]  # Maps from "task-priority", can be keyed-by project or level
     workers: Workers
     taskgraph: TaskGraphConfig
 
     # Optional fields
     docker_image_kind: Optional[str] = None  # Maps from "docker-image-kind"
     task_deadline_after: Optional[Union[str, dict]] = (
-        None  # Maps from "task-deadline-after", can be keyed-by
+        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."""
+        # Validate task_priority can be keyed-by project or level
+        if isinstance(self.task_priority, dict):
+            # Create a validator that accepts TaskPriority values
+            def validate_priority(x):
+                valid_priorities = [
+                    "highest",
+                    "very-high",
+                    "high",
+                    "medium",
+                    "low",
+                    "very-low",
+                    "lowest",
+                ]
+                if x not in valid_priorities:
+                    raise ValueError(f"Invalid task priority: {x}")
+                return x
+
+            validator = optionally_keyed_by("project", "level", validate_priority)
+            # Just validate - it will raise an error if invalid
+            validator(self.task_priority)
+
+        # Validate task_deadline_after can be keyed-by project
+        if self.task_deadline_after and isinstance(self.task_deadline_after, dict):
+            validator = optionally_keyed_by("project", str)
+            # Just validate - it will raise an error if invalid
+            validator(self.task_deadline_after)
+
 
 # Msgspec schema is now the main schema
 graph_config_schema = GraphConfigSchema

src/taskgraph/parameters.py

@@ -175,8 +175,8 @@ def extend_parameters_schema(schema, defaults_fn=None):
     # Store the extension schema for use during validation
     _schema_extensions.append(schema)
 
-    # Schema extension is no longer supported with msgspec.Struct inheritance
-    # Extensions are tracked in _schema_extensions list instead
+    # With msgspec, schema extensions are tracked in the _schema_extensions list
+    # for validation purposes rather than being merged into a single schema
 
     if defaults_fn:
         defaults_functions.append(defaults_fn)
@@ -240,92 +240,39 @@ def _fill_defaults(repo_root=None, **kwargs):
         return kwargs
 
     def check(self):
-        # For msgspec schemas, we need to validate differently
-        if isinstance(base_schema, type) and issubclass(base_schema, msgspec.Struct):
-            try:
-                # Convert underscore keys to kebab-case for msgspec validation
-                params = self.copy()
-                # BaseSchema uses kebab-case (rename="kebab"), so we need to convert keys
-                kebab_params = {}
-                for k, v in params.items():
-                    # Convert underscore to kebab-case
-                    kebab_key = k.replace("_", "-")
-                    kebab_params[kebab_key] = v
-
-                # Handle extensions if present
-                global _schema_extensions
-                for ext_schema in _schema_extensions:
-                    if isinstance(ext_schema, dict):
-                        # Simple dict validation - just check if required keys exist
-                        for key in ext_schema:
-                            # Just skip validation of extensions for now
-                            pass
-
-                if self.strict:
-                    # Strict validation with msgspec
-                    # First check for extra fields
-                    schema_fields = {
-                        f.encode_name for f in msgspec.structs.fields(base_schema)
-                    }
-
-                    # Add extension fields if present
-                    for ext_schema in _schema_extensions:
-                        if isinstance(ext_schema, dict):
-                            for key in ext_schema.keys():
-                                # Extract field name
-                                if hasattr(key, "key"):
-                                    field_name = key.key.replace("_", "-")
-                                else:
-                                    field_name = str(key).replace("_", "-")
-                                schema_fields.add(field_name)
-
-                    extra_fields = set(kebab_params.keys()) - schema_fields
-                    if extra_fields:
-                        raise ParameterMismatch(
-                            f"Invalid parameters: Extra fields not allowed: {extra_fields}"
-                        )
-                    # Now validate the base schema fields
-                    base_fields = {
-                        f.encode_name for f in msgspec.structs.fields(base_schema)
-                    }
-                    base_params = {
-                        k: v for k, v in kebab_params.items() if k in base_fields
-                    }
-                    msgspec.convert(base_params, base_schema)
-                else:
-                    # Non-strict: validate only the fields that exist in the schema
-                    # Filter to only schema fields
-                    schema_fields = {
-                        f.encode_name for f in msgspec.structs.fields(base_schema)
-                    }
-                    filtered_params = {
-                        k: v for k, v in kebab_params.items() if k in schema_fields
-                    }
-                    msgspec.convert(filtered_params, base_schema)
-            except (msgspec.ValidationError, msgspec.DecodeError) as e:
-                raise ParameterMismatch(f"Invalid parameters: {e}")
-        else:
-            # For non-msgspec schemas, validate using the Schema class
-            from taskgraph.util.schema import validate_schema  # noqa: PLC0415
-
-            try:
-                if self.strict:
-                    validate_schema(base_schema, self.copy(), "Invalid parameters:")
-                else:
-                    # In non-strict mode, allow extra fields
-                    if hasattr(base_schema, "allow_extra"):
-                        original_allow_extra = base_schema.allow_extra
-                        base_schema.allow_extra = True
-                        try:
-                            validate_schema(
-                                base_schema, self.copy(), "Invalid parameters:"
-                            )
-                        finally:
-                            base_schema.allow_extra = original_allow_extra
-                    else:
-                        validate_schema(base_schema, self.copy(), "Invalid parameters:")
-            except Exception as e:
-                raise ParameterMismatch(str(e))
+        # Validate parameters using msgspec schema
+        try:
+            # Convert underscore keys to kebab-case since BaseSchema uses rename="kebab"
+            kebab_params = {k.replace("_", "-"): v for k, v in self.items()}
+
+            if self.strict:
+                # Strict mode: validate against schema and check for extra fields
+                # Get all valid field names from the base schema
+                schema_fields = {
+                    f.encode_name for f in msgspec.structs.fields(base_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, base_schema)
+            else:
+                # Non-strict mode: only validate fields that exist in the schema
+                # Filter to only include fields defined in the schema
+                schema_fields = {
+                    f.encode_name for f in msgspec.structs.fields(base_schema)
+                }
+                filtered_params = {
+                    k: v for k, v in kebab_params.items() if k in schema_fields
+                }
+                msgspec.convert(filtered_params, base_schema)
+        except (msgspec.ValidationError, msgspec.DecodeError) as e:
+            raise ParameterMismatch(f"Invalid parameters: {e}")
 
     def __getitem__(self, k):
         try:

src/taskgraph/transforms/docker_image.py

@@ -30,34 +30,27 @@
 
 #: Schema for docker_image transforms
 class DockerImageSchema(Schema):
-    """
-    Schema for docker_image transforms.
-
-    Attributes:
-        name: Name of the docker image.
-        parent: Name of the parent docker image.
-        symbol: Treeherder symbol.
-        task_from: Relative path (from config.path) to the file the docker image was defined in.
-        args: Arguments to use for the Dockerfile.
-        definition: Name of the docker image definition under taskcluster/docker, when
-                   different from the docker image name.
-        packages: List of package tasks this docker image depends on.
-        index: Information for indexing this build so its artifacts can be discovered.
-        cache: Whether this image should be cached based on inputs.
-    """
-
     # Required field first
+    # Name of the docker image.
     name: str
 
     # Optional fields
+    # Name of the parent docker image.
     parent: Optional[str] = None
+    # Treeherder symbol.
     symbol: Optional[str] = None
+    # Relative path (from config.path) to the file the docker image was defined in.
     task_from: Optional[str] = None
+    # Arguments to use for the Dockerfile.
     args: Optional[Dict[str, str]] = None
+    # Name of the docker image definition under taskcluster/docker, when
+    # different from the docker image name.
     definition: Optional[str] = None
+    # List of package tasks this docker image depends on.
     packages: Optional[List[str]] = None
-    # For now, use Any for index since task_description_schema is not converted yet
+    # Information for indexing this build so its artifacts can be discovered.
     index: Optional[Any] = None
+    # Whether this image should be cached based on inputs.
     cache: Optional[bool] = None
 
 

src/taskgraph/transforms/fetch.py

@@ -33,32 +33,29 @@ class FetchConfig(Schema, rename=None, omit_defaults=False):
 
 
 class FetchSchema(Schema):
-    """
-    Schema for fetch transforms.
-
-    Attributes:
-        name: Name of the task.
-        task_from: Relative path (from config.path) to the file the task was defined in.
-        description: Description of the task.
-        expires_after: When the task expires.
-        docker_image: Docker image configuration.
-        fetch_alias: An alias that can be used instead of the real fetch task name in
-                    fetch stanzas for tasks.
-        artifact_prefix: The prefix of the taskcluster artifact being uploaded.
-                        Defaults to `public/`; if it starts with something other than
-                        `public/` the artifact will require scopes to access.
-        attributes: Task attributes.
-        fetch: Fetch configuration with type and additional fields.
-    """
-
+    # Required fields
+    # Name of the task.
     name: str
+    # Description of the task.
     description: str
+    # Fetch configuration with type and additional fields.
     fetch: Dict[str, Any]  # Must have 'type' key, other keys depend on type
+
+    # Optional fields
+    # Relative path (from config.path) to the file the task was defined in.
     task_from: Optional[str] = None
+    # When the task expires.
     expires_after: Optional[str] = None
+    # Docker image configuration.
     docker_image: Optional[Any] = None
+    # An alias that can be used instead of the real fetch task name in
+    # fetch stanzas for tasks.
     fetch_alias: Optional[str] = None
+    # The prefix of the taskcluster artifact being uploaded.
+    # Defaults to `public/`; if it starts with something other than
+    # `public/` the artifact will require scopes to access.
     artifact_prefix: Optional[str] = None
+    # Task attributes.
     attributes: Optional[Dict[str, Any]] = None
 
     def __post_init__(self):

src/taskgraph/transforms/from_deps.py

@@ -24,7 +24,6 @@
 from taskgraph.util.set_name import SET_NAME_MAP
 
 
-# Define FetchEntry for the fetches field
 class FetchEntry(Schema, rename=None):
     """A fetch entry for an artifact."""
 
@@ -33,40 +32,36 @@ class FetchEntry(Schema, rename=None):
 
 
 class FromDepsConfig(Schema):
-    """
-    Configuration for from-deps transforms.
-
-    Attributes:
-        kinds: Limit dependencies to specified kinds (defaults to all kinds in
-              `kind-dependencies`). The first kind in the list is the "primary" kind.
-              The dependency of this kind will be used to derive the label
-              and copy attributes (if `copy-attributes` is True).
-        set_name: UPDATE ME AND DOCS. Can be a string from SET_NAME_MAP, False, None,
-                 or a dict with a SET_NAME_MAP key.
-        with_attributes: Limit dependencies to tasks whose attributes match
-                        using :func:`~taskgraph.util.attributes.attrmatch`.
-        group_by: 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.
-        copy_attributes: If True, copy attributes from the dependency matching the
-                        first kind in the `kinds` list (whether specified explicitly
-                        or taken from `kind-dependencies`).
-        unique_kinds: If true (the default), there must be only a single unique task
-                     for each kind in a dependency group. Setting this to false
-                     disables that requirement.
-        fetches: If present, a `fetches` entry will be added for each task
-                dependency. Attributes of the upstream task may be used as
-                substitution values in the `artifact` or `dest` values of the
-                `fetches` entry.
-    """
-
+    # Optional fields
+    # Limit dependencies to specified kinds (defaults to all kinds in
+    # `kind-dependencies`).
+    #
+    # The first kind in the list is the "primary" kind. The
+    # 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
+    # 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
+    # If True, copy attributes from the dependency matching the
+    # first kind in the `kinds` list (whether specified explicitly
+    # or taken from `kind-dependencies`).
     copy_attributes: Optional[bool] = None
+    # If true (the default), there must be only a single unique task
+    # for each kind in a dependency group. Setting this to false
+    # disables that requirement.
     unique_kinds: Optional[bool] = None
+    # If present, a `fetches` entry will be added for each task
+    # dependency. Attributes of the upstream task may be used as
+    # substitution values in the `artifact` or `dest` values of the
+    # `fetches` entry.
     fetches: Optional[Dict[str, List[Union[str, Dict[str, str]]]]] = None
 
     def __post_init__(self):