feat: update schema from voluptuous to msgspec

Author: abhishekmadan30

docs/concepts/transforms.rst

@@ -105,17 +105,16 @@ about the state of the tasks at given points. Here is an example:
 
 .. code-block:: python
 
-   from voluptuous import Optional, Required
-
+   from typing import Optional
    from taskgraph.transforms.base import TransformSequence
    from taskgraph.util.schema import Schema
 
-   my_schema = Schema({
-       Required("foo"): str,
-       Optional("bar"): bool,
-   })
+   class MySchema(Schema):
+       foo: str  # Required field
+       bar: Optional[bool] = None  # Optional field
 
-   transforms.add_validate(my_schema)
+   transforms = TransformSequence()
+   transforms.add_validate(MySchema)
 
 In the above example, we can be sure that every task dict has a string field
 called ``foo``, and may or may not have a boolean field called ``bar``.

docs/tutorials/creating-a-task-graph.rst

@@ -136,16 +136,16 @@ comments for explanations):
 
 .. code-block:: python
 
-   from voluptuous import Optional, Required
-
-   from taskgraph.transforms.base import TransformSequence
+   from typing import Optional
    from taskgraph.util.schema import Schema
+   from taskgraph.transforms.base import TransformSequence
+
+   # Define the schema using Schema base class.
+   class HelloDescriptionSchema(Schema):
+       text: str                          # Required field
+       description: Optional[str] = None  # Optional field
 
-   # Define the schema. We use the `voluptuous` package to handle validation.
-   hello_description_schema = Schema({
-       Required("text"): str,
-       Optional("description"): str,
-   })
+   hello_description_schema = HelloDescriptionSchema
 
    # Create a 'TransformSequence' instance. This class collects transform
    # functions to run later.

pyproject.toml

@@ -25,13 +25,14 @@ dependencies = [
   "cookiecutter~=2.1",
   "json-e>=2.7",
   "mozilla-repo-urls",
+  "msgspec>=0.18.6",
   "PyYAML>=5.3.1",
   "redo>=2.0",
   "requests>=2.25",
   "slugid>=2.0",
   "taskcluster>=91.0",
   "taskcluster-urls>=11.0",
-  "voluptuous>=0.12.1",
+  "voluptuous>=0.14.2",
 ]
 
 [dependency-groups]

src/taskgraph/config.py

@@ -2,112 +2,121 @@
 # 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
 from dataclasses import dataclass
 from pathlib import Path
-
-from voluptuous import ALLOW_EXTRA, All, Any, Extra, Length, Optional, Required
+from typing import Dict, List, Literal, Optional, Union
 
 from .util.caches import CACHES
 from .util.python_path import find_object
-from .util.schema import Schema, optionally_keyed_by, validate_schema
+from .util.schema import (
+    Struct,
+    TaskPriority,
+    optionally_keyed_by,
+    validate_schema,
+)
 from .util.vcs import get_repository
 from .util.yaml import load_yaml
 
 logger = logging.getLogger(__name__)
 
+# CacheName type for valid cache names
+CacheName = Literal[tuple(CACHES.keys())]
 
-#: Schema for the graph config
-graph_config_schema = Schema(
-    {
-        # The trust-domain for this graph.
-        # (See https://firefox-source-docs.mozilla.org/taskcluster/taskcluster/taskgraph.html#taskgraph-trust-domain)  # noqa
-        Required("trust-domain"): str,
-        Optional(
-            "docker-image-kind",
-            description="Name of the docker image kind (default: docker-image)",
-        ): str,
-        Required("task-priority"): optionally_keyed_by(
-            "project",
-            "level",
-            Any(
-                "highest",
-                "very-high",
-                "high",
-                "medium",
-                "low",
-                "very-low",
-                "lowest",
-            ),
-        ),
-        Optional(
-            "task-deadline-after",
-            description="Default 'deadline' for tasks, in relative date format. "
-            "Eg: '1 week'",
-        ): optionally_keyed_by("project", str),
-        Optional(
-            "task-expires-after",
-            description="Default 'expires-after' for level 1 tasks, in relative date format. "
-            "Eg: '90 days'",
-        ): str,
-        Required("workers"): {
-            Required("aliases"): {
-                str: {
-                    Required("provisioner"): optionally_keyed_by("level", str),
-                    Required("implementation"): str,
-                    Required("os"): str,
-                    Required("worker-type"): optionally_keyed_by("level", str),
-                }
-            },
-        },
-        Required("taskgraph"): {
-            Optional(
-                "register",
-                description="Python function to call to register extensions.",
-            ): str,
-            Optional("decision-parameters"): str,
-            Optional(
-                "cached-task-prefix",
-                description="The taskcluster index prefix to use for caching tasks. "
-                "Defaults to `trust-domain`.",
-            ): str,
-            Optional(
-                "cache-pull-requests",
-                description="Should tasks from pull requests populate the cache",
-            ): bool,
-            Optional(
-                "index-path-regexes",
-                description="Regular expressions matching index paths to be summarized.",
-            ): [str],
-            Optional(
-                "run",
-                description="Configuration related to the 'run' transforms.",
-            ): {
-                Optional(
-                    "use-caches",
-                    description="List of caches to enable, or a boolean to "
-                    "enable/disable all of them.",
-                ): Any(bool, list(CACHES.keys())),
-            },
-            Required("repositories"): All(
-                {
-                    str: {
-                        Required("name"): str,
-                        Optional("project-regex"): str,
-                        Optional("ssh-secret-name"): str,
-                        # FIXME
-                        Extra: str,
-                    }
-                },
-                Length(min=1),
-            ),
-        },
-    },
-    extra=ALLOW_EXTRA,
-)
+
+class WorkerAliasStruct(Struct):
+    """Worker alias configuration."""
+
+    provisioner: optionally_keyed_by("level", str, use_msgspec=True)  # type: ignore
+    implementation: str
+    os: str
+    worker_type: optionally_keyed_by("level", str, use_msgspec=True)  # type: ignore
+
+
+class WorkersStruct(Struct, rename=None):
+    """Workers configuration."""
+
+    aliases: Dict[str, WorkerAliasStruct]
+
+
+class Repository(Struct, 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"
+
+
+class RunConfig(Struct):
+    """Run transforms configuration."""
+
+    # List of caches to enable, or a boolean to enable/disable all of them.
+    use_caches: Optional[Union[bool, List[str]]] = None  # Maps from "use-caches"
+
+    def __post_init__(self):
+        """Validate that cache names are valid."""
+        if isinstance(self.use_caches, list):
+            invalid = set(self.use_caches) - set(CACHES.keys())
+            if invalid:
+                raise ValueError(
+                    f"Invalid cache names: {invalid}. "
+                    f"Valid names are: {list(CACHES.keys())}"
+                )
+
+
+class TaskGraphStruct(Struct):
+    """Taskgraph specific configuration."""
+
+    # Required fields first
+    repositories: Dict[str, Repository]
+
+    # Optional fields
+    # Python function to call to register extensions.
+    register: Optional[str] = None
+    decision_parameters: Optional[str] = None  # Maps from "decision-parameters"
+    # The taskcluster index prefix to use for caching tasks. Defaults to `trust-domain`.
+    cached_task_prefix: Optional[str] = None  # Maps from "cached-task-prefix"
+    # Should tasks from pull requests populate the cache
+    cache_pull_requests: Optional[bool] = None  # Maps from "cache-pull-requests"
+    # Regular expressions matching index paths to be summarized.
+    index_path_regexes: Optional[List[str]] = None  # Maps from "index-path-regexes"
+    # Configuration related to the 'run' transforms.
+    run: Optional[RunConfig] = None
+
+
+class GraphConfigStruct(Struct, forbid_unknown_fields=False):
+    """Main graph configuration schema.
+
+    This schema allows extra fields for flexibility in graph configuration.
+    """
+
+    # Required fields first
+    # The trust-domain for this graph.
+    # (See https://firefox-source-docs.mozilla.org/taskcluster/taskcluster/taskgraph.html#taskgraph-trust-domain)
+    trust_domain: str  # Maps from "trust-domain"
+    task_priority: optionally_keyed_by(
+        "project", "level", TaskPriority, use_msgspec=True
+    )  # type: ignore
+    workers: WorkersStruct
+    taskgraph: TaskGraphStruct
+
+    # Optional fields
+    # Name of the docker image kind (default: docker-image)
+    docker_image_kind: Optional[str] = None  # Maps from "docker-image-kind"
+    # Default 'deadline' for tasks, in relative date format. Eg: '1 week'
+    task_deadline_after: Optional[
+        optionally_keyed_by("project", str, use_msgspec=True)
+    ] = None  # type: ignore
+    # Default 'expires-after' for level 1 tasks, in relative date format. Eg: '90 days'
+    task_expires_after: Optional[str] = None  # Maps from "task-expires-after"
 
 
 @dataclass(frozen=True, eq=False)
@@ -178,7 +187,8 @@ def kinds_dir(self):
 
 
 def validate_graph_config(config):
-    validate_schema(graph_config_schema, config, "Invalid graph configuration:")
+    """Validate graph configuration using msgspec."""
+    validate_schema(GraphConfigStruct, config, "Invalid graph configuration:")
 
 
 def load_graph_config(root_dir):

src/taskgraph/decision.py

@@ -9,9 +9,9 @@
 import shutil
 import time
 from pathlib import Path
+from typing import Any, Dict, Optional
 
 import yaml
-from voluptuous import Optional
 
 from taskgraph.actions import render_actions_json
 from taskgraph.create import create_tasks
@@ -20,8 +20,8 @@
 from taskgraph.taskgraph import TaskGraph
 from taskgraph.util import json
 from taskgraph.util.python_path import find_object
-from taskgraph.util.schema import Schema, validate_schema
-from taskgraph.util.vcs import get_repository
+from taskgraph.util.schema import Struct, validate_schema
+from taskgraph.util.vcs import Repository, get_repository
 from taskgraph.util.yaml import load_yaml
 
 logger = logging.getLogger(__name__)
@@ -40,11 +40,9 @@
 
 
 #: Schema for try_task_config.json version 2
-try_task_config_schema_v2 = Schema(
-    {
-        Optional("parameters"): {str: object},
-    }
-)
+class TryTaskConfigSchemaV2(Struct):
+    # All fields are optional
+    parameters: Optional[Dict[str, Any]] = None
 
 
 def full_task_graph_to_runnable_tasks(full_task_json):
@@ -277,7 +275,7 @@ def set_try_config(parameters, task_config_file):
         task_config_version = task_config.pop("version")
         if task_config_version == 2:
             validate_schema(
-                try_task_config_schema_v2,
+                TryTaskConfigSchemaV2,
                 task_config,
                 "Invalid v2 `try_task_config.json`.",
             )