lsst-felis 28.2024.4500__py3-none-any.whl → 30.0.0rc3__py3-none-any.whl

felis/datamodel.py

@@ -23,29 +23,42 @@
 
 from __future__ import annotations
 
+import json
 import logging
+import sys
 from collections.abc import Sequence
 from enum import StrEnum, auto
-from typing import IO, Annotated, Any, Generic, Literal, TypeAlias, TypeVar, Union
+from typing import IO, Annotated, Any, Generic, Literal, TypeAlias, TypeVar
 
 import yaml
 from astropy import units as units  # type: ignore
 from astropy.io.votable import ucd  # type: ignore
 from lsst.resources import ResourcePath, ResourcePathExpression
-from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator, model_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    PrivateAttr,
+    ValidationError,
+    ValidationInfo,
+    field_serializer,
+    field_validator,
+    model_validator,
+)
+from pydantic_core import InitErrorDetails
 
-from .db.dialects import get_supported_dialects
-from .db.sqltypes import get_type_func
-from .db.utils import string_to_typeengine
+from .db._dialects import get_supported_dialects, string_to_typeengine
+from .db._sqltypes import get_type_func
 from .types import Boolean, Byte, Char, Double, FelisType, Float, Int, Long, Short, String, Text, Unicode
 
 logger = logging.getLogger(__name__)
 
 __all__ = (
     "BaseObject",
-    "Column",
     "CheckConstraint",
+    "Column",
     "Constraint",
+    "DataType",
     "ForeignKeyConstraint",
     "Index",
     "Schema",
@@ -58,6 +71,7 @@ CONFIG = ConfigDict(
     populate_by_name=True,  # Populate attributes by name.
     extra="forbid",  # Do not allow extra fields.
     str_strip_whitespace=True,  # Strip whitespace from string fields.
+    use_enum_values=False,  # Do not use enum values during serialization.
 )
 """Pydantic model configuration as described in:
 https://docs.pydantic.dev/2.0/api/config/#pydantic.config.ConfigDict
@@ -117,7 +131,7 @@ class BaseObject(BaseModel):
 
 
 class DataType(StrEnum):
-    """`Enum` representing the data types supported by Felis."""
+    """``Enum`` representing the data types supported by Felis."""
 
     boolean = auto()
     byte = auto()
@@ -134,6 +148,32 @@ class DataType(StrEnum):
     timestamp = auto()
 
 
+def validate_ivoa_ucd(ivoa_ucd: str) -> str:
+    """Validate IVOA UCD values.
+
+    Parameters
+    ----------
+    ivoa_ucd
+        IVOA UCD value to check.
+
+    Returns
+    -------
+    `str`
+        The IVOA UCD value if it is valid.
+
+    Raises
+    ------
+    ValueError
+        If the IVOA UCD value is invalid.
+    """
+    if ivoa_ucd is not None:
+        try:
+            ucd.parse_ucd(ivoa_ucd, check_controlled_vocabulary=True, has_colon=";" in ivoa_ucd)
+        except ValueError as e:
+            raise ValueError(f"Invalid IVOA UCD: {e}")
+    return ivoa_ucd
+
+
 class Column(BaseObject):
     """Column model."""
 
@@ -159,12 +199,6 @@ class Column(BaseObject):
     autoincrement: bool | None = None
     """Whether the column is autoincremented."""
 
-    mysql_datatype: str | None = Field(None, alias="mysql:datatype")
-    """MySQL datatype override on the column."""
-
-    postgresql_datatype: str | None = Field(None, alias="postgresql:datatype")
-    """PostgreSQL datatype override on the column."""
-
     ivoa_ucd: str | None = Field(None, alias="ivoa:ucd")
     """IVOA UCD of the column."""
 
@@ -193,6 +227,12 @@ class Column(BaseObject):
     votable_datatype: str | None = Field(None, alias="votable:datatype")
     """VOTable datatype of the column."""
 
+    mysql_datatype: str | None = Field(None, alias="mysql:datatype")
+    """MySQL datatype override on the column."""
+
+    postgresql_datatype: str | None = Field(None, alias="postgresql:datatype")
+    """PostgreSQL datatype override on the column."""
+
     @model_validator(mode="after")
     def check_value(self) -> Column:
         """Check that the default value is valid.
@@ -235,12 +275,7 @@ class Column(BaseObject):
         `str`
             The IVOA UCD value if it is valid.
         """
-        if ivoa_ucd is not None:
-            try:
-                ucd.parse_ucd(ivoa_ucd, check_controlled_vocabulary=True, has_colon=";" in ivoa_ucd)
-            except ValueError as e:
-                raise ValueError(f"Invalid IVOA UCD: {e}")
-        return ivoa_ucd
+        return validate_ivoa_ucd(ivoa_ucd)
 
     @model_validator(mode="after")
     def check_units(self) -> Column:
@@ -387,7 +422,7 @@ class Column(BaseObject):
 
     @model_validator(mode="before")
     @classmethod
-    def check_votable_arraysize(cls, values: dict[str, Any]) -> dict[str, Any]:
+    def check_votable_arraysize(cls, values: dict[str, Any], info: ValidationInfo) -> dict[str, Any]:
         """Set the default value for the ``votable_arraysize`` field, which
         corresponds to ``arraysize`` in the IVOA VOTable standard.
 
@@ -395,6 +430,8 @@ class Column(BaseObject):
         ----------
         values
             Values of the column.
+        info
+            Validation context used to determine if the check is enabled.
 
         Returns
         -------
@@ -409,6 +446,7 @@ class Column(BaseObject):
         if values.get("name", None) is None or values.get("datatype", None) is None:
             # Skip bad column data that will not validate
             return values
+        context = info.context if info.context else {}
         arraysize = values.get("votable:arraysize", None)
         if arraysize is None:
             length = values.get("length", None)
@@ -418,7 +456,14 @@ class Column(BaseObject):
                 if datatype == "char":
                     arraysize = str(length)
                 elif datatype in ("string", "unicode", "binary"):
-                    arraysize = f"{length}*"
+                    if context.get("force_unbounded_arraysize", False):
+                        arraysize = "*"
+                        logger.debug(
+                            f"Forced VOTable's 'arraysize' to '*' on column '{values['name']}' with datatype "
+                            + f"'{values['datatype']}' and length '{length}'"
+                        )
+                    else:
+                        arraysize = f"{length}*"
             elif datatype in ("timestamp", "text"):
                 arraysize = "*"
             if arraysize is not None:
@@ -437,6 +482,59 @@ class Column(BaseObject):
                 values["votable:arraysize"] = str(arraysize)
         return values
 
+    @field_serializer("datatype")
+    def serialize_datatype(self, value: DataType) -> str:
+        """Convert `DataType` to string when serializing to JSON/YAML.
+
+        Parameters
+        ----------
+        value
+            The `DataType` value to serialize.
+
+        Returns
+        -------
+        `str`
+            The serialized `DataType` value.
+        """
+        return str(value)
+
+    @field_validator("datatype", mode="before")
+    @classmethod
+    def deserialize_datatype(cls, value: str) -> DataType:
+        """Convert string back into `DataType` when loading from JSON/YAML.
+
+        Parameters
+        ----------
+        value
+            The string value to deserialize.
+
+        Returns
+        -------
+        `DataType`
+            The deserialized `DataType` value.
+        """
+        return DataType(value)
+
+    @model_validator(mode="after")
+    def check_votable_xtype(self) -> Column:
+        """Set the default value for the ``votable_xtype`` field, which
+        corresponds to an Extended Datatype or ``xtype`` in the IVOA VOTable
+        standard.
+
+        Returns
+        -------
+        `Column`
+            The column being validated.
+
+        Notes
+        -----
+        This is currently only set automatically for the Felis ``timestamp``
+        datatype.
+        """
+        if self.datatype == DataType.timestamp and self.votable_xtype is None:
+            self.votable_xtype = "timestamp"
+        return self
+
 
 class Constraint(BaseObject):
     """Table constraint model."""
@@ -472,6 +570,22 @@ class CheckConstraint(Constraint):
     expression: str
     """Expression for the check constraint."""
 
+    @field_serializer("type")
+    def serialize_type(self, value: str) -> str:
+        """Ensure '@type' is included in serialized output.
+
+        Parameters
+        ----------
+        value
+            The value to serialize.
+
+        Returns
+        -------
+        `str`
+            The serialized value.
+        """
+        return value
+
 
 class UniqueConstraint(Constraint):
     """Table unique constraint model."""
@@ -482,12 +596,30 @@ class UniqueConstraint(Constraint):
     columns: list[str]
     """Columns in the unique constraint."""
 
+    @field_serializer("type")
+    def serialize_type(self, value: str) -> str:
+        """Ensure '@type' is included in serialized output.
+
+        Parameters
+        ----------
+        value
+            The value to serialize.
+
+        Returns
+        -------
+        `str`
+            The serialized value.
+        """
+        return value
+
 
 class ForeignKeyConstraint(Constraint):
     """Table foreign key constraint model.
 
     This constraint is used to define a foreign key relationship between two
-    tables in the schema.
+    tables in the schema. There must be at least one column in the
+    `columns` list, and at least one column in the `referenced_columns` list
+    or a validation error will be raised.
 
     Notes
     -----
@@ -498,12 +630,62 @@ class ForeignKeyConstraint(Constraint):
     type: Literal["ForeignKey"] = Field("ForeignKey", alias="@type")
     """Type of the constraint."""
 
-    columns: list[str]
+    columns: list[str] = Field(min_length=1)
     """The columns comprising the foreign key."""
 
-    referenced_columns: list[str] = Field(alias="referencedColumns")
+    referenced_columns: list[str] = Field(alias="referencedColumns", min_length=1)
     """The columns referenced by the foreign key."""
 
+    on_delete: Literal["CASCADE", "SET NULL", "SET DEFAULT", "RESTRICT", "NO ACTION"] | None = None
+    """Action to take when the referenced row is deleted."""
+
+    on_update: Literal["CASCADE", "SET NULL", "SET DEFAULT", "RESTRICT", "NO ACTION"] | None = None
+    """Action to take when the referenced row is updated."""
+
+    @field_serializer("type")
+    def serialize_type(self, value: str) -> str:
+        """Ensure '@type' is included in serialized output.
+
+        Parameters
+        ----------
+        value
+            The value to serialize.
+
+        Returns
+        -------
+        `str`
+            The serialized value.
+        """
+        return value
+
+    @model_validator(mode="after")
+    def check_column_lengths(self) -> ForeignKeyConstraint:
+        """Check that the `columns` and `referenced_columns` lists have the
+        same length.
+
+        Returns
+        -------
+        `ForeignKeyConstraint`
+            The foreign key constraint being validated.
+
+        Raises
+        ------
+        ValueError
+            Raised if the `columns` and `referenced_columns` lists do not have
+            the same length.
+        """
+        if len(self.columns) != len(self.referenced_columns):
+            raise ValueError(
+                "Columns and referencedColumns must have the same length for a ForeignKey constraint"
+            )
+        return self
+
+
+_ConstraintType = Annotated[
+    CheckConstraint | ForeignKeyConstraint | UniqueConstraint, Field(discriminator="type")
+]
+"""Type alias for a constraint type."""
+
 
 class Index(BaseObject):
     """Table index model.
@@ -545,23 +727,91 @@ class Index(BaseObject):
         return values
 
 
-_ConstraintType = Annotated[
-    Union[CheckConstraint, ForeignKeyConstraint, UniqueConstraint], Field(discriminator="type")
-]
-"""Type alias for a constraint type."""
+ColumnRef: TypeAlias = str
+"""Type alias for a column reference."""
 
 
-class Table(BaseObject):
-    """Table model."""
+class ColumnGroup(BaseObject):
+    """Column group model."""
 
-    columns: Sequence[Column]
-    """Columns in the table."""
+    columns: list[ColumnRef | Column] = Field(..., min_length=1)
+    """Columns in the group."""
 
-    constraints: list[_ConstraintType] = Field(default_factory=list)
-    """Constraints on the table."""
+    ivoa_ucd: str | None = Field(None, alias="ivoa:ucd")
+    """IVOA UCD of the column."""
 
-    indexes: list[Index] = Field(default_factory=list)
-    """Indexes on the table."""
+    table: Table | None = Field(None, exclude=True)
+    """Reference to the parent table."""
+
+    @field_validator("ivoa_ucd")
+    @classmethod
+    def check_ivoa_ucd(cls, ivoa_ucd: str) -> str:
+        """Check that IVOA UCD values are valid.
+
+        Parameters
+        ----------
+        ivoa_ucd
+            IVOA UCD value to check.
+
+        Returns
+        -------
+        `str`
+            The IVOA UCD value if it is valid.
+        """
+        return validate_ivoa_ucd(ivoa_ucd)
+
+    @model_validator(mode="after")
+    def check_unique_columns(self) -> ColumnGroup:
+        """Check that the columns list contains unique items.
+
+        Returns
+        -------
+        `ColumnGroup`
+            The column group being validated.
+        """
+        column_ids = [col if isinstance(col, str) else col.id for col in self.columns]
+        if len(column_ids) != len(set(column_ids)):
+            raise ValueError("Columns in the group must be unique")
+        return self
+
+    def _dereference_columns(self) -> None:
+        """Dereference ColumnRef to Column objects."""
+        if self.table is None:
+            raise ValueError("ColumnGroup must have a reference to its parent table")
+
+        dereferenced_columns: list[ColumnRef | Column] = []
+        for col in self.columns:
+            if isinstance(col, str):
+                # Dereference ColumnRef to Column object
+                try:
+                    col_obj = self.table._find_column_by_id(col)
+                except KeyError as e:
+                    raise ValueError(f"Column '{col}' not found in table '{self.table.name}'") from e
+                dereferenced_columns.append(col_obj)
+            else:
+                dereferenced_columns.append(col)
+
+        self.columns = dereferenced_columns
+
+    @field_serializer("columns")
+    def serialize_columns(self, columns: list[ColumnRef | Column]) -> list[str]:
+        """Serialize columns as their IDs.
+
+        Parameters
+        ----------
+        columns
+            The columns to serialize.
+
+        Returns
+        -------
+        `list` [ `str` ]
+            The serialized column IDs.
+        """
+        return [col if isinstance(col, str) else col.id for col in columns]
+
+
+class Table(BaseObject):
+    """Table model."""
 
     primary_key: str | list[str] | None = Field(None, alias="primaryKey")
     """Primary key of the table."""
@@ -575,6 +825,18 @@ class Table(BaseObject):
     mysql_charset: str | None = Field(None, alias="mysql:charset")
     """MySQL charset to use for the table."""
 
+    columns: Sequence[Column]
+    """Columns in the table."""
+
+    column_groups: list[ColumnGroup] = Field(default_factory=list, alias="columnGroups")
+    """Column groups in the table."""
+
+    constraints: list[_ConstraintType] = Field(default_factory=list)
+    """Constraints on the table."""
+
+    indexes: list[Index] = Field(default_factory=list)
+    """Indexes on the table."""
+
     @field_validator("columns", mode="after")
     @classmethod
     def check_unique_column_names(cls, columns: list[Column]) -> list[Column]:
@@ -653,6 +915,43 @@ class Table(BaseObject):
                 return self
         raise ValueError(f"Table '{self.name}' is missing at least one column designated as 'tap:principal'")
 
+    def _find_column_by_id(self, id: str) -> Column:
+        """Find a column by ID.
+
+        Parameters
+        ----------
+        id
+            The ID of the column to find.
+
+        Returns
+        -------
+        `Column`
+            The column with the given ID.
+
+        Raises
+        ------
+        ValueError
+            Raised if the column is not found.
+        """
+        for column in self.columns:
+            if column.id == id:
+                return column
+        raise KeyError(f"Column '{id}' not found in table '{self.name}'")
+
+    @model_validator(mode="after")
+    def dereference_column_groups(self: Table) -> Table:
+        """Dereference columns in column groups.
+
+        Returns
+        -------
+        `Table`
+            The table with dereferenced column groups.
+        """
+        for group in self.column_groups:
+            group.table = self
+            group._dereference_columns()
+        return self
+
 
 class SchemaVersion(BaseModel):
     """Schema version model."""
@@ -696,10 +995,10 @@ class SchemaIdVisitor:
         if hasattr(obj, "id"):
             obj_id = getattr(obj, "id")
             if self.schema is not None:
-                if obj_id in self.schema.id_map:
+                if obj_id in self.schema._id_map:
                     self.duplicates.add(obj_id)
                 else:
-                    self.schema.id_map[obj_id] = obj
+                    self.schema._id_map[obj_id] = obj
 
     def visit_schema(self, schema: Schema) -> None:
         """Visit the objects in a schema and build the ID map.
@@ -757,6 +1056,56 @@ class SchemaIdVisitor:
 T = TypeVar("T", bound=BaseObject)
 
 
+def _strip_ids(data: Any) -> Any:
+    """Recursively strip '@id' fields from a dictionary or list.
+
+    Parameters
+    ----------
+    data
+        The data to strip IDs from, which can be a dictionary, list, or any
+        other type. Other types will be returned unchanged.
+    """
+    if isinstance(data, dict):
+        data.pop("@id", None)
+        for k, v in data.items():
+            data[k] = _strip_ids(v)
+        return data
+    elif isinstance(data, list):
+        return [_strip_ids(item) for item in data]
+    else:
+        return data
+
+
+def _append_error(
+    errors: list[InitErrorDetails],
+    loc: tuple,
+    input_value: Any,
+    error_message: str,
+    error_type: str = "value_error",
+) -> None:
+    """Append an error to the errors list.
+
+    Parameters
+    ----------
+    errors : list[InitErrorDetails]
+        The list of errors to append to.
+    loc : tuple
+        The location of the error in the schema.
+    input_value : Any
+        The input value that caused the error.
+    error_message : str
+        The error message to include in the context.
+    """
+    errors.append(
+        {
+            "type": error_type,
+            "loc": loc,
+            "input": input_value,
+            "ctx": {"error": error_message},
+        }
+    )
+
+
 class Schema(BaseObject, Generic[T]):
     """Database schema model.
 
@@ -769,7 +1118,7 @@ class Schema(BaseObject, Generic[T]):
     tables: Sequence[Table]
     """The tables in the schema."""
 
-    id_map: dict[str, Any] = Field(default_factory=dict, exclude=True)
+    _id_map: dict[str, Any] = PrivateAttr(default_factory=dict)
     """Map of IDs to objects."""
 
     @model_validator(mode="before")
@@ -807,6 +1156,14 @@ class Schema(BaseObject, Generic[T]):
                         if "@id" not in column:
                             column["@id"] = f"#{table['name']}.{column['name']}"
                             logger.debug(f"Generated ID '{column['@id']}' for column '{column['name']}'")
+                if "columnGroups" in table:
+                    for column_group in table["columnGroups"]:
+                        if "@id" not in column_group:
+                            column_group["@id"] = f"#{table['name']}.{column_group['name']}"
+                            logger.debug(
+                                f"Generated ID '{column_group['@id']}' for column group "
+                                f"'{column_group['name']}'"
+                            )
                 if "constraints" in table:
                     for constraint in table["constraints"]:
                         if "@id" not in constraint:
@@ -931,20 +1288,21 @@ class Schema(BaseObject, Generic[T]):
 
         return self
 
-    def _create_id_map(self: Schema) -> Schema:
+    @model_validator(mode="after")
+    def create_id_map(self: Schema) -> Schema:
         """Create a map of IDs to objects.
 
+        Returns
+        -------
+        `Schema`
+            The schema with the ID map created.
+
         Raises
         ------
         ValueError
             Raised if duplicate identifiers are found in the schema.
-
-        Notes
-        -----
-        This is called automatically by the `model_post_init` method. If the
-        ID map is already populated, this method will return immediately.
         """
-        if len(self.id_map):
+        if self._id_map:
             logger.debug("Ignoring call to create_id_map() - ID map was already populated")
             return self
         visitor: SchemaIdVisitor = SchemaIdVisitor()
@@ -953,25 +1311,152 @@ class Schema(BaseObject, Generic[T]):
             raise ValueError(
                 "Duplicate IDs found in schema:\n    " + "\n    ".join(visitor.duplicates) + "\n"
             )
+        logger.debug("Created ID map with %d entries", len(self._id_map))
         return self
 
-    def model_post_init(self, ctx: Any) -> None:
-        """Post-initialization hook for the model.
+    def _validate_column_id(
+        self: Schema,
+        column_id: str,
+        loc: tuple,
+        errors: list[InitErrorDetails],
+    ) -> None:
+        """Validate a column ID from a constraint and append errors if invalid.
 
         Parameters
         ----------
-        ctx
-            The context object which was passed to the model.
+        schema : Schema
+            The schema being validated.
+        column_id : str
+            The column ID to validate.
+        loc : tuple
+            The location of the error in the schema.
+        errors : list[InitErrorDetails]
+            The list of errors to append to.
+        """
+        if column_id not in self:
+            _append_error(
+                errors,
+                loc,
+                column_id,
+                f"Column ID '{column_id}' not found in schema",
+            )
+        elif not isinstance(self[column_id], Column):
+            _append_error(
+                errors,
+                loc,
+                column_id,
+                f"ID '{column_id}' does not refer to a Column object",
+            )
 
-        Notes
-        -----
-        This method is called automatically by Pydantic after the model is
-        initialized. It is used to create the ID map for the schema.
+    def _validate_foreign_key_column(
+        self: Schema,
+        column_id: str,
+        table: Table,
+        loc: tuple,
+        errors: list[InitErrorDetails],
+    ) -> None:
+        """Validate a foreign key column ID from a constraint and append errors
+        if invalid.
 
-        The ``ctx`` argument has the type `Any` because this is the function
-        signature in Pydantic itself.
+        Parameters
+        ----------
+        schema : Schema
+            The schema being validated.
+        column_id : str
+            The foreign key column ID to validate.
+        loc : tuple
+            The location of the error in the schema.
+        errors : list[InitErrorDetails]
+            The list of errors to append to.
         """
-        self._create_id_map()
+        try:
+            table._find_column_by_id(column_id)
+        except KeyError:
+            _append_error(
+                errors,
+                loc,
+                column_id,
+                f"Column '{column_id}' not found in table '{table.name}'",
+            )
+
+    @model_validator(mode="after")
+    def check_constraints(self: Schema) -> Schema:
+        """Check constraint objects for validity. This needs to be deferred
+        until after the schema is fully loaded and the ID map is created.
+
+        Raises
+        ------
+        pydantic.ValidationError
+            Raised if any constraints are invalid.
+
+        Returns
+        -------
+        `Schema`
+            The schema being validated.
+        """
+        errors: list[InitErrorDetails] = []
+
+        for table_index, table in enumerate(self.tables):
+            for constraint_index, constraint in enumerate(table.constraints):
+                column_ids: list[str] = []
+                referenced_column_ids: list[str] = []
+
+                if isinstance(constraint, ForeignKeyConstraint):
+                    column_ids += constraint.columns
+                    referenced_column_ids += constraint.referenced_columns
+                elif isinstance(constraint, UniqueConstraint):
+                    column_ids += constraint.columns
+                # No extra checks are required on CheckConstraint objects.
+
+                # Validate the foreign key columns
+                for column_id in column_ids:
+                    self._validate_column_id(
+                        column_id,
+                        (
+                            "tables",
+                            table_index,
+                            "constraints",
+                            constraint_index,
+                            "columns",
+                            column_id,
+                        ),
+                        errors,
+                    )
+                    # Check that the foreign key column is within the source
+                    # table.
+                    self._validate_foreign_key_column(
+                        column_id,
+                        table,
+                        (
+                            "tables",
+                            table_index,
+                            "constraints",
+                            constraint_index,
+                            "columns",
+                            column_id,
+                        ),
+                        errors,
+                    )
+
+                # Validate the primary key (reference) columns
+                for referenced_column_id in referenced_column_ids:
+                    self._validate_column_id(
+                        referenced_column_id,
+                        (
+                            "tables",
+                            table_index,
+                            "constraints",
+                            constraint_index,
+                            "referenced_columns",
+                            referenced_column_id,
+                        ),
+                        errors,
+                    )
+
+        if errors:
+            raise ValidationError.from_exception_data("Schema validation failed", errors)
+
+        return self
 
     def __getitem__(self, id: str) -> BaseObject:
         """Get an object by its ID.
@@ -988,7 +1473,7 @@ class Schema(BaseObject, Generic[T]):
         """
         if id not in self:
             raise KeyError(f"Object with ID '{id}' not found in schema")
-        return self.id_map[id]
+        return self._id_map[id]
 
     def __contains__(self, id: str) -> bool:
         """Check if an object with the given ID is in the schema.
@@ -998,7 +1483,7 @@ class Schema(BaseObject, Generic[T]):
         id
             The ID of the object to check.
         """
-        return id in self.id_map
+        return id in self._id_map
 
     def find_object_by_id(self, id: str, obj_type: type[T]) -> T:
         """Find an object with the given type by its ID.
@@ -1114,3 +1599,58 @@ class Schema(BaseObject, Generic[T]):
         logger.debug("Loading schema from: '%s'", source)
         yaml_data = yaml.safe_load(source)
         return Schema.model_validate(yaml_data, context=context)
+
+    def _model_dump(self, strip_ids: bool = False) -> dict[str, Any]:
+        """Dump the schema as a dictionary with some default arguments
+        applied.
+
+        Parameters
+        ----------
+        strip_ids
+            Whether to strip the IDs from the dumped data. Defaults to `False`.
+
+        Returns
+        -------
+        `dict` [ `str`, `Any` ]
+            The dumped schema data as a dictionary.
+        """
+        data = self.model_dump(by_alias=True, exclude_none=True, exclude_defaults=True)
+        if strip_ids:
+            data = _strip_ids(data)
+        return data
+
+    def dump_yaml(self, stream: IO[str] = sys.stdout, strip_ids: bool = False) -> None:
+        """Pretty print the schema as YAML.
+
+        Parameters
+        ----------
+        stream
+            The stream to write the YAML data to.
+        strip_ids
+            Whether to strip the IDs from the dumped data. Defaults to `False`.
+        """
+        data = self._model_dump(strip_ids=strip_ids)
+        yaml.safe_dump(
+            data,
+            stream,
+            default_flow_style=False,
+            sort_keys=False,
+        )
+
+    def dump_json(self, stream: IO[str] = sys.stdout, strip_ids: bool = False) -> None:
+        """Pretty print the schema as JSON.
+
+        Parameters
+        ----------
+        stream
+            The stream to write the JSON data to.
+        strip_ids
+            Whether to strip the IDs from the dumped data. Defaults to `False`.
+        """
+        data = self._model_dump(strip_ids=strip_ids)
+        json.dump(
+            data,
+            stream,
+            indent=4,
+            sort_keys=False,
+        )