feat: used schema wrapper for msgspec

Author: abhishekmadan30

src/taskgraph/config.py

@@ -13,7 +13,7 @@
 import msgspec
 
 from .util.python_path import find_object
-from .util.schema import validate_schema
+from .util.schema import Schema, validate_schema
 from .util.vcs import get_repository
 from .util.yaml import load_yaml
 
@@ -26,7 +26,7 @@
 ]
 
 
-class WorkerAlias(msgspec.Struct, kw_only=True, rename="kebab"):
+class WorkerAlias(Schema):
     """Worker alias configuration."""
 
     provisioner: Union[str, dict]
@@ -35,32 +35,38 @@ class WorkerAlias(msgspec.Struct, kw_only=True, rename="kebab"):
     worker_type: Union[str, dict]  # Can be keyed-by, maps from "worker-type"
 
 
-class Workers(msgspec.Struct, kw_only=True):
+class Workers(Schema, rename=None):
     """Workers configuration."""
 
     aliases: Dict[str, WorkerAlias]
 
 
-class Repository(msgspec.Struct, kw_only=True, rename="kebab"):
+class Repository(Schema):
     """Repository 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(msgspec.Struct, kw_only=True, rename="kebab"):
+class RunConfig(Schema):
     """Run transforms configuration."""
 
     use_caches: Optional[Union[bool, List[str]]] = None  # Maps from "use-caches"
 
 
-class TaskGraphConfig(msgspec.Struct, kw_only=True, rename="kebab"):
+class TaskGraphConfig(Schema):
     """Taskgraph specific configuration."""
 
+    # Required fields first
     repositories: Dict[str, Repository]
+
+    # Optional fields
     register: Optional[str] = None
     decision_parameters: Optional[str] = None  # Maps from "decision-parameters"
     cached_task_prefix: Optional[str] = None  # Maps from "cached-task-prefix"
@@ -69,17 +75,18 @@ class TaskGraphConfig(msgspec.Struct, kw_only=True, rename="kebab"):
     run: Optional[RunConfig] = None
 
 
-class GraphConfigSchema(
-    msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"
-):
+class GraphConfigSchema(Schema):
     """Main graph configuration schema."""
 
+    # Required fields first
     trust_domain: str  # Maps from "trust-domain"
     task_priority: Union[
         TaskPriority, dict
     ]  # Maps from "task-priority", can be keyed-by
     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
@@ -161,9 +168,7 @@ 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:", use_msgspec=True
-    )
+    validate_schema(GraphConfigSchema, config, "Invalid graph configuration:")
 
 
 def load_graph_config(root_dir):

src/taskgraph/decision.py

@@ -11,7 +11,6 @@
 from pathlib import Path
 from typing import Any, Dict, Optional
 
-import msgspec
 import yaml
 
 from taskgraph.actions import render_actions_json
@@ -21,7 +20,7 @@
 from taskgraph.taskgraph import TaskGraph
 from taskgraph.util import json
 from taskgraph.util.python_path import find_object
-from taskgraph.util.schema import validate_schema
+from taskgraph.util.schema import Schema, validate_schema
 from taskgraph.util.vcs import Repository, get_repository
 from taskgraph.util.yaml import load_yaml
 
@@ -41,7 +40,8 @@
 
 
 #: Schema for try_task_config.json version 2
-class TryTaskConfigSchemaV2(msgspec.Struct, kw_only=True, omit_defaults=True):
+class TryTaskConfigSchemaV2(Schema):
+    # All fields are optional
     parameters: Optional[Dict[str, Any]] = None
 
 
@@ -358,7 +358,6 @@ def set_try_config(parameters, task_config_file):
                 try_task_config_schema_v2,
                 task_config,
                 "Invalid v2 `try_task_config.json`.",
-                use_msgspec=True,
             )
             parameters.update(task_config["parameters"])
             return

src/taskgraph/parameters.py

@@ -20,6 +20,7 @@
 
 from taskgraph.util import json, yaml
 from taskgraph.util.readonlydict import ReadOnlyDict
+from taskgraph.util.schema import Schema
 from taskgraph.util.taskcluster import find_task_id, get_artifact_url
 from taskgraph.util.vcs import get_repository
 
@@ -28,20 +29,22 @@ class ParameterMismatch(Exception):
     """Raised when a parameters.yml has extra or missing parameters."""
 
 
-class CodeReviewConfig(msgspec.Struct, kw_only=True, rename="kebab"):
+class CodeReviewConfig(Schema):
     """Code review configuration."""
 
+    # Required field
     phabricator_build_target: str
 
 
 #: Schema for base parameters.
 #: Please keep this list sorted and in sync with docs/reference/parameters.rst
-class BaseSchema(msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"):
+class BaseSchema(Schema):
     """Base parameters schema.
 
     This defines the core parameters that all taskgraph runs require.
     """
 
+    # Required fields (most are required)
     base_repository: str
     base_ref: str
     base_rev: str
@@ -58,8 +61,6 @@ class BaseSchema(msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab
     head_tag: str
     level: str
     moz_build_date: str
-    next_version: Optional[str]
-    optimize_strategies: Optional[str]
     optimize_target_tasks: bool
     owner: str
     project: str
@@ -70,6 +71,10 @@ class BaseSchema(msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab
     # used at run-time
     target_tasks_method: str
     tasks_for: str
+
+    # Optional fields
+    next_version: Optional[str]
+    optimize_strategies: Optional[str]
     version: Optional[str]
     code_review: Optional[CodeReviewConfig] = None
 
@@ -157,7 +162,7 @@ def extend_parameters_schema(schema, defaults_fn=None):
     graph-configuration.
 
     Args:
-        schema: The schema object (dict or msgspec) used to describe extended
+        schema: The schema object (msgspec) used to describe extended
             parameters.
         defaults_fn (function): A function which takes no arguments and returns a
             dict mapping parameter name to default value in the
@@ -170,9 +175,8 @@ def extend_parameters_schema(schema, defaults_fn=None):
     # Store the extension schema for use during validation
     _schema_extensions.append(schema)
 
-    # Also extend the base_schema if it's a Schema instance
-    if hasattr(base_schema, "extend"):
-        base_schema = base_schema.extend(schema)
+    # Schema extension is no longer supported with msgspec.Struct inheritance
+    # Extensions are tracked in _schema_extensions list instead
 
     if defaults_fn:
         defaults_functions.append(defaults_fn)

src/taskgraph/transforms/base.py

@@ -8,8 +8,6 @@
 from dataclasses import dataclass, field
 from typing import Dict, List, Union
 
-import msgspec
-
 from taskgraph.task import Task
 
 from ..config import GraphConfig
@@ -156,9 +154,5 @@ def __call__(self, config, tasks):
                 )
             else:
                 error = "In unknown task:"
-            # Check if schema is a msgspec.Struct type
-            use_msgspec = isinstance(self.schema, type) and issubclass(
-                self.schema, msgspec.Struct
-            )
-            validate_schema(self.schema, task, error, use_msgspec=use_msgspec)
+            validate_schema(self.schema, task, error)
             yield task

src/taskgraph/transforms/run/__init__.py

@@ -29,7 +29,7 @@
 
 
 # Fetches schema using msgspec
-class FetchesSchema(msgspec.Struct, kw_only=True, rename="kebab", omit_defaults=True):
+class FetchesSchema(Schema):
     """Schema for fetch configuration."""
 
     artifact: str
@@ -39,14 +39,14 @@ class FetchesSchema(msgspec.Struct, kw_only=True, rename="kebab", omit_defaults=
 
 
 # When configuration using msgspec
-class WhenConfig(msgspec.Struct, kw_only=True, rename="kebab", omit_defaults=True):
+class WhenConfig(Schema):
     """Configuration for when a task should be included."""
 
     files_changed: List[str] = msgspec.field(default_factory=list)
 
 
 # Run configuration using msgspec
-class RunConfig(msgspec.Struct, kw_only=True, omit_defaults=True):
+class RunConfig(Schema, rename=None):
     """Configuration for how to run a task."""
 
     using: str
@@ -56,20 +56,19 @@ class RunConfig(msgspec.Struct, kw_only=True, omit_defaults=True):
 
 
 # Run description schema using msgspec
-class RunDescriptionSchema(
-    msgspec.Struct, kw_only=True, rename="kebab", omit_defaults=True
-):
+class RunDescriptionSchema(Schema):
     """Schema for run transforms."""
 
-    # Task naming
-    name: TOptional[str] = None
-    label: TOptional[str] = None
-
-    # Required fields
+    # Required fields first
     description: str
     run: RunConfig
     worker_type: str
 
+    # Optional fields
+    # Task naming
+    name: TOptional[str] = None
+    label: TOptional[str] = None
+
     # Optional fields from task description
     priority: TOptional[str] = None
     attributes: Dict[str, Any] = msgspec.field(default_factory=dict)
@@ -104,7 +103,7 @@ class RunDescriptionSchema(
 fetches_schema = FetchesSchema
 
 #: Schema for a run transforms - now using msgspec
-run_description_schema = Schema(RunDescriptionSchema)
+run_description_schema = RunDescriptionSchema
 
 
 transforms = TransformSequence()
@@ -399,15 +398,13 @@ def wrap(func):
 
 
 # Simple schema for always-optimized
-class AlwaysOptimizedRunSchema(msgspec.Struct, kw_only=True):
+class AlwaysOptimizedRunSchema(Schema, omit_defaults=False):
     """Schema for always-optimized run tasks."""
 
     using: str = "always-optimized"
 
 
-@run_task_using(
-    "always-optimized", "always-optimized", Schema(AlwaysOptimizedRunSchema)
-)
+@run_task_using("always-optimized", "always-optimized", AlwaysOptimizedRunSchema)
 def always_optimized(config, task, taskdesc):
     pass
 

src/taskgraph/transforms/run/index_search.py

@@ -10,16 +10,15 @@
 
 from typing import List
 
-import msgspec
-
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.transforms.run import run_task_using
+from taskgraph.util.schema import Schema
 
 transforms = TransformSequence()
 
 
 #: Schema for run.using index-search
-class RunTaskSchema(msgspec.Struct, kw_only=True, rename="kebab"):
+class RunTaskSchema(Schema):
     using: str
     # A list of indexes in decreasing order of priority at which to lookup for this
     # task. This is interpolated with the graph parameters.

src/taskgraph/transforms/run/run_task.py

@@ -9,8 +9,6 @@
 import os
 from typing import Dict, List, Literal, Optional, Union
 
-import msgspec
-
 from taskgraph.transforms.run import run_task_using
 from taskgraph.transforms.run.common import (
     support_caches,
@@ -26,14 +24,24 @@
 
 
 #: Schema for run.using run_task
-class RunTaskSchema(msgspec.Struct, kw_only=True, rename="kebab", omit_defaults=True):
+class RunTaskSchema(Schema):
     """
     Schema for run.using run_task.
     """
 
+    # Required fields first
     # Specifies the task type. Must be 'run-task'.
     using: Literal["run-task"]
 
+    # The command arguments to pass to the `run-task` script, after the checkout
+    # arguments. If a list, it will be passed directly; otherwise it will be
+    # included in a single argument to the command specified by `exec-with`.
+    command: Union[List[Union[str, Dict[str, str]]], str, Dict[str, str]]
+
+    # Base work directory used to set up the task.
+    workdir: str
+
+    # Optional fields
     # Specifies which caches to use. May take a boolean in which case either all
     # (True) or no (False) caches will be used. Alternatively, it can accept a
     # list of caches to enable. Defaults to only the checkout cache enabled.
@@ -51,11 +59,6 @@ class RunTaskSchema(msgspec.Struct, kw_only=True, rename="kebab", omit_defaults=
     # directory where sparse profiles are defined (build/sparse-profiles/).
     sparse_profile: Optional[str] = None
 
-    # The command arguments to pass to the `run-task` script, after the checkout
-    # arguments. If a list, it will be passed directly; otherwise it will be
-    # included in a single argument to the command specified by `exec-with`.
-    command: Union[List[Union[str, Dict[str, str]]], str, Dict[str, str]]
-
     # Specifies what to execute the command with in the event the command is a
     # string.
     exec_with: Optional[Literal["bash", "powershell"]] = None
@@ -64,14 +67,11 @@ class RunTaskSchema(msgspec.Struct, kw_only=True, rename="kebab", omit_defaults=
     # or Python installation is in a non-standard location on the workers.
     run_task_command: Optional[List[str]] = None
 
-    # Base work directory used to set up the task.
-    workdir: str
-
     # Whether to run as root. Defaults to False.
     run_as_root: bool = False
 
 
-run_task_schema = Schema(RunTaskSchema)
+run_task_schema = RunTaskSchema
 
 
 def common_setup(config, task, taskdesc, command):