feat: convert parameter schema from voluptuous to msgspec

Author: abhishekmadan30

src/taskgraph/parameters.py

@@ -10,16 +10,17 @@
 from io import BytesIO
 from pprint import pformat
 from subprocess import CalledProcessError
+from typing import Optional, Union
 from unittest.mock import Mock
 from urllib.parse import urlparse
 from urllib.request import urlopen
 
 import mozilla_repo_urls
-from voluptuous import ALLOW_EXTRA, Any, Optional, Required, Schema
+import msgspec
 
 from taskgraph.util import json, yaml
 from taskgraph.util.readonlydict import ReadOnlyDict
-from taskgraph.util.schema import validate_schema
+from taskgraph.util.schema import Schema, validate_schema
 from taskgraph.util.taskcluster import find_task_id, get_artifact_url
 from taskgraph.util.vcs import get_repository
 
@@ -28,44 +29,52 @@ class ParameterMismatch(Exception):
     """Raised when a parameters.yml has extra or missing parameters."""
 
 
+class CodeReviewConfig(Schema, kw_only=True):
+    phabricator_build_target: str
+
+
 #: Schema for base parameters.
 #: Please keep this list sorted and in sync with docs/reference/parameters.rst
-base_schema = Schema(
-    {
-        Required("base_repository"): str,
-        Optional("base_ref"): str,
-        Required("base_rev"): str,
-        Required("build_date"): int,
-        Required("build_number"): int,
-        Required("do_not_optimize"): [str],
-        Required("enable_always_target"): Any(bool, [str]),
-        Required("existing_tasks"): {str: str},
-        Required("files_changed"): [str],
-        Required("filters"): [str],
-        Required("head_ref"): str,
-        Required("head_repository"): str,
-        Required("head_rev"): str,
-        Required("head_tag"): str,
-        Required("level"): str,
-        Required("moz_build_date"): str,
-        Required("next_version"): Any(str, None),
-        Required("optimize_strategies"): Any(str, None),
-        Required("optimize_target_tasks"): bool,
-        Required("owner"): str,
-        Required("project"): str,
-        Required("pushdate"): int,
-        Required("pushlog_id"): str,
-        Required("repository_type"): str,
-        # target-kinds is not included, since it should never be
-        # used at run-time
-        Required("target_tasks_method"): str,
-        Required("tasks_for"): str,
-        Required("version"): Any(str, None),
-        Optional("code-review"): {
-            Required("phabricator-build-target"): str,
-        },
-    }
-)
+class BaseParametersSchema(
+    Schema,
+    kw_only=True,
+    forbid_unknown_fields=False,
+    rename={"code_review": "code-review"},
+):
+    base_repository: str
+    base_rev: str
+    build_date: int
+    build_number: int
+    do_not_optimize: list[str]
+    enable_always_target: Union[bool, list[str]]
+    existing_tasks: dict[str, str]
+    files_changed: list[str]
+    filters: list[str]
+    head_ref: str
+    head_repository: str
+    head_rev: str
+    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
+    pushdate: int
+    pushlog_id: str
+    repository_type: str
+    # target-kinds is not included, since it should never be
+    # used at run-time
+    target_tasks_method: str
+    tasks_for: str
+    version: Optional[str]
+    base_ref: Optional[str] = None
+    code_review: Optional[CodeReviewConfig] = None
+
+
+base_schema = BaseParametersSchema
+_parameter_extensions: list = []
 
 
 def get_contents(path):
@@ -143,19 +152,27 @@ def extend_parameters_schema(schema, defaults_fn=None):
     graph-configuration.
 
     Args:
-        schema (Schema): The voluptuous.Schema object used to describe extended
-            parameters.
+        schema: A msgspec ``Schema`` subclass describing extended parameters.
         defaults_fn (function): A function which takes no arguments and returns a
             dict mapping parameter name to default value in the
             event strict=False (optional).
     """
-    global base_schema
     global defaults_functions
-    base_schema = base_schema.extend(schema)
+    if not (isinstance(schema, type) and issubclass(schema, msgspec.Struct)):
+        raise TypeError(
+            "extend_parameters_schema requires a msgspec Schema subclass; "
+            f"got {type(schema).__name__}"
+        )
+    _parameter_extensions.append(schema)
     if defaults_fn:
         defaults_functions.append(defaults_fn)
 
 
+def _schema_key_names(schema) -> set:
+    """Return the data-level field names declared by a parameters schema."""
+    return {f.encode_name for f in msgspec.structs.fields(schema)}
+
+
 class Parameters(ReadOnlyDict):
     """An immutable dictionary with nicer KeyError messages on failure"""
 
@@ -214,11 +231,30 @@ def _fill_defaults(repo_root=None, **kwargs):
         return kwargs
 
     def check(self):
-        schema = (
-            base_schema if self.strict else base_schema.extend({}, extra=ALLOW_EXTRA)
-        )
+        data = dict(self.copy())
         try:
-            validate_schema(schema, self.copy(), "Invalid parameters:")
+            # Validate core fields against just the subset of data owned by the
+            # base schema. Extension keys are validated separately below, and a
+            # strict-mode check rejects anything unknown to either.
+            base_keys = _schema_key_names(base_schema)
+            base_data = {k: v for k, v in data.items() if k in base_keys}
+            validate_schema(base_schema, base_data, "Invalid parameters:")
+
+            # Validate each registered extension against the keys it declares.
+            allowed = set(base_keys)
+            for ext in _parameter_extensions:
+                ext_keys = _schema_key_names(ext)
+                allowed |= ext_keys
+                ext_data = {k: data[k] for k in ext_keys if k in data}
+                validate_schema(ext, ext_data, "Invalid parameters:")
+
+            # Strict mode: reject any data key not covered by base or extensions.
+            if self.strict:
+                unknown = sorted(set(data) - allowed)
+                if unknown:
+                    raise Exception(
+                        "Invalid parameters:\nunknown keys: " + ", ".join(unknown)
+                    )
         except Exception as e:
             raise ParameterMismatch(str(e))
 

taskcluster/self_taskgraph/custom_parameters.py

@@ -3,10 +3,10 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 import os
-
-from voluptuous import All, Any, Range, Required
+from typing import Optional
 
 from taskgraph.parameters import extend_parameters_schema
+from taskgraph.util.schema import Schema
 
 
 def get_defaults(repo_root):
@@ -15,12 +15,18 @@ def get_defaults(repo_root):
     }
 
 
-extend_parameters_schema(
-    {
-        Required("pull_request_number"): Any(All(int, Range(min=1)), None),
-    },
-    defaults_fn=get_defaults,
-)
+class CustomParametersSchema(Schema, kw_only=True, rename=None):
+    pull_request_number: Optional[int]
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.pull_request_number is not None and self.pull_request_number < 1:
+            raise ValueError(
+                f"pull_request_number must be >= 1, got {self.pull_request_number}"
+            )
+
+
+extend_parameters_schema(CustomParametersSchema, defaults_fn=get_defaults)
 
 
 def decision_parameters(graph_config, parameters):

test/test_parameters.py

@@ -7,11 +7,11 @@
 import gzip
 import os
 from base64 import b64decode
+from typing import Optional
 from unittest import TestCase, mock
 
 import mozilla_repo_urls
 import pytest
-from voluptuous import Optional, Required, Schema
 
 import taskgraph  # noqa: F401
 from taskgraph import parameters
@@ -21,6 +21,7 @@
     extend_parameters_schema,
     load_parameters_file,
 )
+from taskgraph.util.schema import Schema
 
 from .mockedopen import MockedOpen
 
@@ -274,20 +275,19 @@ def test_parameters_format_spec(spec, expected):
 
 
 def test_extend_parameters_schema(monkeypatch):
-    monkeypatch.setattr(
-        parameters,
-        "base_schema",
-        Schema(
-            {
-                Required("foo"): str,
-            }
-        ),
-    )
+    class FooSchema(Schema, kw_only=True, rename=None):
+        foo: str
+
+    class BarSchema(Schema, kw_only=True, rename=None):
+        bar: Optional[bool] = None
+
+    monkeypatch.setattr(parameters, "base_schema", FooSchema)
     monkeypatch.setattr(
         parameters,
         "defaults_functions",
         list(parameters.defaults_functions),
     )
+    monkeypatch.setattr(parameters, "_parameter_extensions", [])
 
     with pytest.raises(ParameterMismatch):
         Parameters(strict=False).check()
@@ -296,9 +296,7 @@ def test_extend_parameters_schema(monkeypatch):
         Parameters(foo="1", bar=True).check()
 
     extend_parameters_schema(
-        {
-            Optional("bar"): bool,
-        },
+        BarSchema,
         defaults_fn=lambda root: {"foo": "1", "bar": False},
     )