feat: cleaned up schema wrapper

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
 
-   import msgspec
    from typing import Optional
    from taskgraph.transforms.base import TransformSequence
    from taskgraph.util.schema import Schema
 
-   class MySchema(msgspec.Struct, kw_only=True):
+   class MySchema(Schema):
        foo: str                   # Required field
        bar: Optional[bool] = None # Optional field
 
    transforms = TransformSequence()
-   transforms.add_validate(Schema(MySchema))
+   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,17 +136,16 @@ comments for explanations):
 
 .. code-block:: python
 
-   import msgspec
    from typing import Optional
    from taskgraph.util.schema import Schema
    from taskgraph.transforms.base import TransformSequence
 
-   # Define the schema using msgspec for better type checking and performance.
-   class HelloDescriptionSchema(msgspec.Struct, kw_only=True):
+   # Define the schema using Schema base class for better type checking and performance.
+   class HelloDescriptionSchema(Schema):
        text: str                          # Required field
        description: Optional[str] = None  # Optional field
 
-   hello_description_schema = Schema(HelloDescriptionSchema)
+   hello_description_schema = HelloDescriptionSchema
 
    # Create a 'TransformSequence' instance. This class collects transform
    # functions to run later.

src/taskgraph/transforms/chunking.py

@@ -7,10 +7,11 @@
 import msgspec
 
 from taskgraph.transforms.base import TransformSequence
+from taskgraph.util.schema import Schema
 from taskgraph.util.templates import substitute
 
 
-class ChunkConfig(msgspec.Struct, kw_only=True, rename="kebab"):
+class ChunkConfig(Schema):
     """
     `chunk` can be used to split one task into `total-chunks`
     tasks, substituting `this_chunk` and `total_chunks` into any
@@ -25,7 +26,7 @@ class ChunkConfig(msgspec.Struct, kw_only=True, rename="kebab"):
 
 
 #: Schema for chunking transforms
-class ChunkSchema(msgspec.Struct, kw_only=True, omit_defaults=True):
+class ChunkSchema(Schema):
     # Optional, so it can be used for a subset of tasks in a kind
     chunk: Optional[ChunkConfig] = None
     __extras__: Dict[str, Any] = msgspec.field(default_factory=dict)

src/taskgraph/transforms/docker_image.py

@@ -7,12 +7,11 @@
 import re
 from typing import Any, Dict, List, Optional
 
-import msgspec
-
 import taskgraph
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.util import json
 from taskgraph.util.docker import create_context_tar, generate_context_hash
+from taskgraph.util.schema import Schema
 
 logger = logging.getLogger(__name__)
 
@@ -30,9 +29,7 @@
 
 
 #: Schema for docker_image transforms
-class DockerImageSchema(
-    msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"
-):
+class DockerImageSchema(Schema):
     """
     Schema for docker_image transforms.
 
@@ -49,7 +46,10 @@ class DockerImageSchema(
         cache: Whether this image should be cached based on inputs.
     """
 
+    # Required field first
     name: str
+
+    # Optional fields
     parent: Optional[str] = None
     symbol: Optional[str] = None
     task_from: Optional[str] = None

src/taskgraph/transforms/fetch.py

@@ -17,21 +17,22 @@
 
 from ..util import path
 from ..util.cached_tasks import add_optimization
+from ..util.schema import Schema
 from ..util.treeherder import join_symbol
 from .base import TransformSequence
 
 CACHE_TYPE = "content.v1"
 
 
 #: Schema for fetch transforms
-class FetchConfig(msgspec.Struct, kw_only=True):
+class FetchConfig(Schema, rename=None, omit_defaults=False):
     """Configuration for a fetch task type."""
 
     type: str
     # Additional fields handled dynamically by fetch builders
 
 
-class FetchSchema(msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"):
+class FetchSchema(Schema):
     """
     Schema for fetch transforms.
 
@@ -199,7 +200,7 @@ def make_task(config, tasks):
         yield task_desc
 
 
-class GPGSignatureConfig(msgspec.Struct, kw_only=True, rename="kebab"):
+class GPGSignatureConfig(Schema):
     """GPG signature verification configuration."""
 
     # URL where GPG signature document can be obtained. Can contain the
@@ -209,9 +210,7 @@ class GPGSignatureConfig(msgspec.Struct, kw_only=True, rename="kebab"):
     key_path: str
 
 
-class StaticUrlFetchConfig(
-    msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"
-):
+class StaticUrlFetchConfig(Schema, rename="kebab"):
     """Configuration for static-url fetch type."""
 
     type: str
@@ -303,7 +302,7 @@ def create_fetch_url_task(config, name, fetch):
     }
 
 
-class GitFetchConfig(msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"):
+class GitFetchConfig(Schema):
     """Configuration for git fetch type."""
 
     type: str

src/taskgraph/transforms/from_deps.py

@@ -20,19 +20,19 @@
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.util.attributes import attrmatch
 from taskgraph.util.dependencies import GROUP_BY_MAP, get_dependencies
-from taskgraph.util.schema import validate_schema
+from taskgraph.util.schema import Schema, validate_schema
 from taskgraph.util.set_name import SET_NAME_MAP
 
 
 # Define FetchEntry for the fetches field
-class FetchEntry(msgspec.Struct, kw_only=True, omit_defaults=True):
+class FetchEntry(Schema, rename=None):
     """A fetch entry for an artifact."""
 
     artifact: str
     dest: Optional[str] = None
 
 
-class FromDepsConfig(msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"):
+class FromDepsConfig(Schema):
     """
     Configuration for from-deps transforms.
 
@@ -94,7 +94,7 @@ def __post_init__(self):
 
 
 #: Schema for from_deps transforms
-class FromDepsSchema(msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"):
+class FromDepsSchema(Schema):
     """Schema for from_deps transforms."""
 
     from_deps: FromDepsConfig

src/taskgraph/transforms/matrix.py

@@ -13,10 +13,11 @@
 import msgspec
 
 from taskgraph.transforms.base import TransformSequence
+from taskgraph.util.schema import Schema
 from taskgraph.util.templates import substitute_task_fields
 
 
-class MatrixConfig(msgspec.Struct, kw_only=True, rename="kebab"):
+class MatrixConfig(Schema):
     """
     Matrix configuration for generating multiple tasks.
     """
@@ -43,7 +44,7 @@ class MatrixConfig(msgspec.Struct, kw_only=True, rename="kebab"):
 
 
 #: Schema for matrix transforms
-class MatrixSchema(msgspec.Struct, kw_only=True, omit_defaults=True):
+class MatrixSchema(Schema):
     name: str
     matrix: Optional[MatrixConfig] = None
     __extras__: Dict[str, Any] = msgspec.field(default_factory=dict)

src/taskgraph/transforms/notify.py

@@ -13,7 +13,7 @@
 import msgspec
 
 from taskgraph.transforms.base import TransformSequence
-from taskgraph.util.schema import resolve_keyed_by
+from taskgraph.util.schema import Schema, resolve_keyed_by
 
 StatusType = Literal[
     "on-completed",
@@ -26,31 +26,31 @@
 ]
 
 
-class EmailRecipient(msgspec.Struct, kw_only=True, rename="kebab"):
+class EmailRecipient(Schema):
     """Email notification recipient."""
 
     type: Literal["email"]
     address: Union[str, Dict[str, Any]]  # Can be keyed-by
     status_type: Optional[StatusType] = None
 
 
-class MatrixRoomRecipient(msgspec.Struct, kw_only=True, rename="kebab"):
+class MatrixRoomRecipient(Schema):
     """Matrix room notification recipient."""
 
     type: Literal["matrix-room"]
     room_id: str
     status_type: Optional[StatusType] = None
 
 
-class PulseRecipient(msgspec.Struct, kw_only=True, rename="kebab"):
+class PulseRecipient(Schema):
     """Pulse notification recipient."""
 
     type: Literal["pulse"]
     routing_key: str
     status_type: Optional[StatusType] = None
 
 
-class SlackChannelRecipient(msgspec.Struct, kw_only=True, rename="kebab"):
+class SlackChannelRecipient(Schema):
     """Slack channel notification recipient."""
 
     type: Literal["slack-channel"]
@@ -71,22 +71,22 @@ class SlackChannelRecipient(msgspec.Struct, kw_only=True, rename="kebab"):
 """Map each type to its primary key that will be used in the route."""
 
 
-class EmailLink(msgspec.Struct, kw_only=True):
+class EmailLink(Schema, rename=None, omit_defaults=False):
     """Email link configuration."""
 
     text: str
     href: str
 
 
-class EmailContent(msgspec.Struct, kw_only=True, omit_defaults=True):
+class EmailContent(Schema, rename=None):
     """Email notification content."""
 
     subject: Optional[str] = None
     content: Optional[str] = None
     link: Optional[EmailLink] = None
 
 
-class MatrixContent(msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"):
+class MatrixContent(Schema):
     """Matrix notification content."""
 
     body: Optional[str] = None
@@ -95,32 +95,30 @@ class MatrixContent(msgspec.Struct, kw_only=True, omit_defaults=True, rename="ke
     msg_type: Optional[str] = None
 
 
-class SlackContent(msgspec.Struct, kw_only=True, omit_defaults=True):
+class SlackContent(Schema, rename=None):
     """Slack notification content."""
 
     text: Optional[str] = None
     blocks: Optional[List[Any]] = None
     attachments: Optional[List[Any]] = None
 
 
-class NotifyContent(msgspec.Struct, kw_only=True, omit_defaults=True):
+class NotifyContent(Schema, rename=None):
     """Notification content configuration."""
 
     email: Optional[EmailContent] = None
     matrix: Optional[MatrixContent] = None
     slack: Optional[SlackContent] = None
 
 
-class NotifyConfig(msgspec.Struct, kw_only=True, omit_defaults=True):
+class NotifyConfig(Schema, rename=None):
     """Modern notification configuration."""
 
     recipients: List[Dict[str, Any]]  # Will be validated as Recipient union
     content: Optional[NotifyContent] = None
 
 
-class LegacyNotificationsConfig(
-    msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"
-):
+class LegacyNotificationsConfig(Schema, rename="kebab"):
     """Legacy notification configuration for backwards compatibility."""
 
     emails: Union[List[str], Dict[str, Any]]  # Can be keyed-by
@@ -130,9 +128,7 @@ class LegacyNotificationsConfig(
 
 
 #: Schema for notify transforms
-class NotifySchema(
-    msgspec.Struct, kw_only=True, omit_defaults=True, tag_field="notify_type"
-):
+class NotifySchema(Schema, tag_field="notify_type"):
     """Schema for notify transforms.
 
     Note: This schema allows either 'notify' or 'notifications' field,

src/taskgraph/transforms/run/toolchain.py

@@ -18,15 +18,14 @@
 )
 from taskgraph.util import path as mozpath
 from taskgraph.util.hash import hash_paths
+from taskgraph.util.schema import Schema
 from taskgraph.util.shell import quote as shell_quote
 
 CACHE_TYPE = "toolchains.v3"
 
 
 #: Schema for run.using toolchain
-class ToolchainRunSchema(
-    msgspec.Struct, kw_only=True, omit_defaults=True, rename="kebab"
-):
+class ToolchainRunSchema(Schema):
     """
     Schema for toolchain-script run configuration.
 
@@ -49,11 +48,14 @@ class ToolchainRunSchema(
         workdir: Base work directory used to set up the task.
     """
 
+    # Required fields first
     using: Literal["toolchain-script"]
     script: str
     sparse_profile: Optional[str]  # Can be None to skip sparse profile
     toolchain_artifact: str
     workdir: str
+
+    # Optional fields
     arguments: Optional[List[str]] = None
     resources: Optional[List[str]] = None
     toolchain_alias: Optional[Union[str, List[str]]] = None