fix: updated optionally keyed by fucntionality

Type refactor

Author: abhishekmadan30

src/taskgraph/config.py

@@ -2,7 +2,6 @@
 # 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 logging
 import os
 import sys
@@ -26,25 +25,10 @@
 
 class WorkerAlias(Schema):
     """Worker alias configuration."""
-
-    provisioner: Union[str, dict]  # Can be keyed-by level
+    provisioner: optionally_keyed_by("level", str)  # type: ignore
     implementation: str
     os: str
-    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)
+    worker_type: optionally_keyed_by("level", str)  # type: ignore
 
 
 class Workers(Schema, rename=None):
@@ -96,47 +80,15 @@ class GraphConfigSchema(Schema, forbid_unknown_fields=False):
 
     # Required fields first
     trust_domain: str  # Maps from "trust-domain"
-    task_priority: Union[
-        TaskPriority, dict
-    ]  # Maps from "task-priority", can be keyed-by project or level
+    task_priority: optionally_keyed_by("project", "level", TaskPriority)  # type: ignore
     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 project
-    )
+    task_deadline_after: Optional[optionally_keyed_by("project", str)] = None  # type: ignore
     task_expires_after: Optional[str] = None  # Maps from "task-expires-after"
 
-    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

src/taskgraph/util/schema.py

@@ -2,9 +2,9 @@
 # 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 pprint
-from typing import List
+from functools import reduce
+from typing import Dict, List, Literal, Union
 
 import msgspec
 
@@ -37,72 +37,18 @@ def validate_schema(schema, obj, msg_prefix):
         raise Exception(f"{msg_prefix}\n{str(exc)}\n{pprint.pformat(obj)}")
 
 
-def optionally_keyed_by(*arguments):
+def UnionTypes(*types):
+    """Use `functools.reduce` to simulate `Union[*allowed_types]` on older
+    Python versions.
     """
-    Mark a schema value as optionally keyed by any of a number of fields.  The
-    schema is the last argument, and the remaining fields are taken to be the
-    field names.  For example:
+    return reduce(lambda a, b: Union[a, b], types)
 
-        'some-value': optionally_keyed_by(
-            'test-platform', 'build-platform',
-            Any('a', 'b', 'c'))
 
-    The resulting schema will allow nesting of `by-test-platform` and
-    `by-build-platform` in either order.
-    """
-    schema = arguments[-1]
+def optionally_keyed_by(*arguments):
+    _type = arguments[-1]
     fields = arguments[:-1]
-
-    def validator(obj):
-        if isinstance(obj, dict) and len(obj) == 1:
-            k, v = list(obj.items())[0]
-            if k.startswith("by-") and k[len("by-") :] in fields:
-                res = {}
-                for kk, vv in v.items():
-                    try:
-                        res[kk] = validator(vv)
-                    except Exception as e:
-                        raise ValueError(f"Error in {k}.{kk}: {str(e)}") from e
-                return res
-            elif k.startswith("by-"):
-                # Unknown by-field
-                raise ValueError(f"Unknown key {k}")
-        # Validate against the schema
-        if isinstance(schema, type) and issubclass(schema, Schema):
-            return schema.validate(obj)
-        elif schema is str:
-            # String validation
-            if not isinstance(obj, str):
-                raise TypeError(f"Expected string, got {type(obj).__name__}")
-            return obj
-        elif schema is int:
-            # Int validation
-            if not isinstance(obj, int):
-                raise TypeError(f"Expected int, got {type(obj).__name__}")
-            return obj
-        elif isinstance(schema, type):
-            # Type validation for built-in types
-            if not isinstance(obj, schema):
-                raise TypeError(f"Expected {schema.__name__}, got {type(obj).__name__}")
-            return obj
-        elif callable(schema):
-            # Other callable validators
-            try:
-                return schema(obj)
-            except:
-                raise
-        else:
-            # Simple type validation
-            if not isinstance(obj, schema):
-                raise TypeError(
-                    f"Expected {getattr(schema, '__name__', str(schema))}, got {type(obj).__name__}"
-                )
-            return obj
-
-    # set to assist autodoc
-    setattr(validator, "schema", schema)
-    setattr(validator, "fields", fields)
-    return validator
+    bykeys = [Literal[f"by-{field}"] for field in fields]
+    return Union[_type, Dict[UnionTypes(*bykeys), Dict[str, _type]]]
 
 
 def resolve_keyed_by(

taskcluster/config.yml

@@ -12,7 +12,10 @@ index:
     products:
         - taskgraph
 
-task-priority: low
+task-priority:
+    by-level:
+        "3": medium
+        default: low
 
 taskgraph:
     register: self_taskgraph:register