Support defining msgspec schemas as dicts

src/taskgraph/transforms/chunking.py

@@ -2,30 +2,29 @@
 # 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 copy
-from typing import Optional
 
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.util.schema import Schema
 from taskgraph.util.templates import substitute
 
-
-class ChunkConfig(Schema):
-    # The total number of chunks to split the task into.
-    total_chunks: int
-    # A list of fields that need to have `{this_chunk}` and/or
-    # `{total_chunks}` replaced in them.
-    substitution_fields: list[str] = []
-
-
-#: Schema for chunking transforms
-class ChunkSchema(Schema, forbid_unknown_fields=False, kw_only=True):
-    # `chunk` can be used to split one task into `total-chunks`
-    # tasks, substituting `this_chunk` and `total_chunks` into any
-    # fields in `substitution-fields`.
-    chunk: Optional[ChunkConfig] = None
-
-
-CHUNK_SCHEMA = ChunkSchema
+CHUNK_SCHEMA = Schema.from_dict(
+    {
+        # `chunk` can be used to split one task into `total-chunks`
+        # tasks, substituting `this_chunk` and `total_chunks` into any
+        # fields in `substitution-fields`.
+        "chunk": Schema.from_dict(
+            {
+                # The total number of chunks to split the task into.
+                "total-chunks": int,
+                # A list of fields that need to have `{this_chunk}` and/or
+                # `{total_chunks}` replaced in them.
+                "substitution-fields": (list[str], []),
+            },
+            optional=True,
+        ),
+    },
+    forbid_unknown_fields=False,
+)
 
 transforms = TransformSequence()
 transforms.add_validate(CHUNK_SCHEMA)

src/taskgraph/transforms/docker_image.py

@@ -5,7 +5,6 @@
 import logging
 import os
 import re
-from typing import Optional
 
 import taskgraph
 from taskgraph.transforms.base import TransformSequence
@@ -27,31 +26,32 @@
 
 transforms = TransformSequence()
 
-
 #: Schema for docker_image transforms
-class DockerImageSchema(Schema):
-    # Name of the docker image.
-    name: str
-    # 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
-    # Information for indexing this build so its artifacts can be discovered.
-    index: Optional[IndexSchema] = None
-    # Whether this image should be cached based on inputs.
-    cache: Optional[bool] = None
-
-
-docker_image_schema = DockerImageSchema
+DOCKER_IMAGE_SCHEMA = Schema.from_dict(
+    {
+        # Name of the docker image.
+        "name": str,
+        # Name of the parent docker image.
+        "parent": (str, None),
+        # Treeherder symbol.
+        "symbol": (str, None),
+        # Relative path (from config.path) to the file the docker image was defined in.
+        "task-from": (str, None),
+        # Arguments to use for the Dockerfile.
+        "args": (dict[str, str], None),
+        # Name of the docker image definition under taskcluster/docker, when
+        # different from the docker image name.
+        "definition": (str, None),
+        # List of package tasks this docker image depends on.
+        "packages": (list[str], None),
+        # Information for indexing this build so its artifacts can be discovered.
+        "index": (IndexSchema, None),
+        # Whether this image should be cached based on inputs.
+        "cache": (bool, None),
+    },
+)
+
+docker_image_schema = DOCKER_IMAGE_SCHEMA
 
 
 transforms.add_validate(docker_image_schema)

src/taskgraph/transforms/fetch.py

@@ -21,36 +21,38 @@
 
 CACHE_TYPE = "content.v1"
 
-
-class FetchSubSchema(Schema, forbid_unknown_fields=False, kw_only=True):
-    # The fetch type
-    type: str
-
+_FETCH_SUB_SCHEMA = Schema.from_dict(
+    {
+        # The fetch type
+        "type": str,
+    },
+    forbid_unknown_fields=False,
+)
 
 #: Schema for fetch transforms
-class FetchSchema(Schema):
-    # Name of the task.
-    name: str
-    # Description of the task.
-    description: str
-    # The fetch configuration
-    fetch: FetchSubSchema
-    # Relative path (from config.path) to the file the task was defined
-    # in.
-    task_from: Optional[str] = None
-    expires_after: Optional[str] = None
-    docker_image: Optional[object] = 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
-    attributes: Optional[dict[str, object]] = None
-
-
-FETCH_SCHEMA = FetchSchema
+FETCH_SCHEMA = Schema.from_dict(
+    {
+        # Name of the task.
+        "name": str,
+        # Description of the task.
+        "description": str,
+        # The fetch configuration
+        "fetch": _FETCH_SUB_SCHEMA,
+        # Relative path (from config.path) to the file the task was defined
+        # in.
+        "task-from": (str, None),
+        "expires-after": (str, None),
+        "docker-image": (object, None),
+        # An alias that can be used instead of the real fetch task name in
+        # fetch stanzas for tasks.
+        "fetch-alias": (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": (str, None),
+        "attributes": (dict[str, object], None),
+    },
+)
 
 # define a collection of payload builders, depending on the worker implementation
 fetch_builders = {}
@@ -173,42 +175,48 @@ def make_task(config, tasks):
         yield task_desc
 
 
-class GpgSignatureConfig(Schema):
-    # URL where GPG signature document can be obtained. Can contain the
-    # value ``{url}``, which will be substituted with the value from
-    # ``url``.
-    sig_url: str
-    # Path to file containing GPG public key(s) used to validate
-    # download.
-    key_path: str
-
-
-class StaticUrlFetchSchema(Schema, forbid_unknown_fields=False, kw_only=True):
-    type: Literal["static-url"]
-    # The URL to download.
-    url: str
-    # The SHA-256 of the downloaded content.
-    sha256: str
-    # Size of the downloaded entity, in bytes.
-    size: int
-    # GPG signature verification.
-    gpg_signature: Optional[GpgSignatureConfig] = None
-    # The name to give to the generated artifact. Defaults to the file
-    # portion of the URL. Using a different extension converts the
-    # archive to the given type. Only conversion to .tar.zst is
-    # supported.
-    artifact_name: Optional[str] = None
-    # Strip the given number of path components at the beginning of
-    # each file entry in the archive.
-    # Requires an artifact-name ending with .tar.zst.
-    strip_components: Optional[int] = None
-    # Add the given prefix to each file entry in the archive.
-    # Requires an artifact-name ending with .tar.zst.
-    add_prefix: Optional[str] = None
-    # Headers to pass alongside the request.
-    headers: Optional[dict[str, str]] = None
-    # IMPORTANT: when adding anything that changes the behavior of the task,
-    # it is important to update the digest data used to compute cache hits.
+_GPG_SIGNATURE_SCHEMA = Schema.from_dict(
+    {
+        # URL where GPG signature document can be obtained. Can contain the
+        # value ``{url}``, which will be substituted with the value from
+        # ``url``.
+        "sig-url": str,
+        # Path to file containing GPG public key(s) used to validate
+        # download.
+        "key-path": str,
+    },
+)
+
+StaticUrlFetchSchema = Schema.from_dict(
+    {
+        "type": Literal["static-url"],
+        # The URL to download.
+        "url": str,
+        # The SHA-256 of the downloaded content.
+        "sha256": str,
+        # Size of the downloaded entity, in bytes.
+        "size": int,
+        # GPG signature verification.
+        "gpg-signature": (_GPG_SIGNATURE_SCHEMA, None),
+        # The name to give to the generated artifact. Defaults to the file
+        # portion of the URL. Using a different extension converts the
+        # archive to the given type. Only conversion to .tar.zst is
+        # supported.
+        "artifact-name": (str, None),
+        # Strip the given number of path components at the beginning of
+        # each file entry in the archive.
+        # Requires an artifact-name ending with .tar.zst.
+        "strip-components": (int, None),
+        # Add the given prefix to each file entry in the archive.
+        # Requires an artifact-name ending with .tar.zst.
+        "add-prefix": (str, None),
+        # Headers to pass alongside the request.
+        "headers": (dict[str, str], None),
+        # IMPORTANT: when adding anything that changes the behavior of the task,
+        # it is important to update the digest data used to compute cache hits.
+    },
+    forbid_unknown_fields=False,
+)
 
 
 @fetch_builder("static-url", schema=StaticUrlFetchSchema)
@@ -274,18 +282,22 @@ def create_fetch_url_task(config, name, fetch):
     }
 
 
-class GitFetchSchema(Schema, forbid_unknown_fields=False, kw_only=True):
-    type: Literal["git"]
-    repo: str
-    revision: str
-    include_dot_git: Optional[bool] = None
-    artifact_name: Optional[str] = None
-    path_prefix: Optional[str] = None
-    # ssh-key is a taskcluster secret path (e.g. project/civet/github-deploy-key)
-    # In the secret dictionary, the key should be specified as
-    #  "ssh_privkey": "-----BEGIN OPENSSH PRIVATE KEY-----\nkfksnb3jc..."
-    # n.b. The OpenSSH private key file format requires a newline at the end of the file.
-    ssh_key: Optional[str] = None
+GitFetchSchema = Schema.from_dict(
+    {
+        "type": Literal["git"],
+        "repo": str,
+        "revision": str,
+        "include-dot-git": (bool, None),
+        "artifact-name": (str, None),
+        "path-prefix": (str, None),
+        # ssh-key is a taskcluster secret path (e.g. project/civet/github-deploy-key)
+        # In the secret dictionary, the key should be specified as
+        #  "ssh_privkey": "-----BEGIN OPENSSH PRIVATE KEY-----\nkfksnb3jc..."
+        # n.b. The OpenSSH private key file format requires a newline at the end of the file.
+        "ssh-key": (str, None),
+    },
+    forbid_unknown_fields=False,
+)
 
 
 @fetch_builder("git", schema=GitFetchSchema)