refactor(RFC): Add demo for PEP487

Ripped from vega@9a48448

altair/utils/schemapi.py

@@ -887,6 +887,118 @@ def _deep_copy(obj: _CopyImpl | Any, by_ref: set[str]) -> _CopyImpl | Any:
         return obj
 
 
+class _SchemaBasePEP487:
+    """Minimal demo for testing feasibility of `__init_subclass__`."""
+
+    _schema: ClassVar[dict[str, Any]]
+    _rootschema: ClassVar[dict[str, Any]]
+    _class_is_valid_at_instantiation: ClassVar[bool] = True
+
+    def __init__(self, *args: Any, **kwds: Any) -> None:
+        if (kwds and args) or len(args) > 1:
+            name = type(self).__name__
+            _args = ", ".join(f"{a!r}" for a in args)
+            _kwds = ", ".join(f"{k}={v!r}" for k, v in kwds.items())
+            msg = (
+                f"Expected either:\n"
+                f" - a single arg with no kwds, for, e.g. {{'type': 'string'}}\n"
+                f" - zero args with zero or more kwds for {{'type': 'object'}}\n\n"
+                f"but got: {name}({_args}, {_kwds})"
+            )
+            raise AssertionError(msg)
+        # use object.__setattr__ because we override setattr below.
+        self._args: tuple[Any, ...]
+        self._kwds: dict[str, Any]
+        object.__setattr__(self, "_args", args)
+        object.__setattr__(self, "_kwds", kwds)
+
+    def __init_subclass__(
+        cls,
+        *args: Any,
+        schema: dict[str, Any] | None = None,
+        rootschema: dict[str, Any] | None = None,
+        valid_at_init: bool | None = None,
+        **kwds: Any,
+    ) -> None:
+        super().__init_subclass__(*args, **kwds)
+        if schema is None:
+            if hasattr(cls, "_schema"):
+                schema = cls._schema
+            else:
+                msg = (
+                    f"Cannot instantiate object of type {cls}: "
+                    "_schema class attribute is not defined."
+                )
+                raise TypeError(msg)
+        if rootschema is None:
+            if hasattr(cls, "_rootschema"):
+                rootschema = cls._rootschema
+            elif "$ref" not in schema:
+                rootschema = schema
+            else:
+                msg = "`rootschema` must be provided if `schema` contains a `'$ref'` and does not inherit one."
+                raise TypeError(msg)
+        if valid_at_init is None:
+            valid_at_init = cls._class_is_valid_at_instantiation
+        cls._schema = schema
+        cls._rootschema = rootschema
+        cls._class_is_valid_at_instantiation = valid_at_init
+
+    @overload
+    def _get(self, attr: str, default: Optional = ...) -> Any | UndefinedType: ...
+    @overload
+    def _get(self, attr: str, default: T) -> Any | T: ...
+    def _get(self, attr: str, default: Optional[T] = Undefined) -> Any | T:
+        """Get an attribute, returning default if not present."""
+        if (item := self._kwds.get(attr, Undefined)) is not Undefined:
+            return item
+        else:
+            return default
+
+    def __dir__(self) -> list[str]:
+        return sorted(chain(super().__dir__(), self._kwds))
+
+    def __eq__(self, other: Any) -> bool:
+        return (
+            type(self) is type(other)
+            and self._args == other._args
+            and self._kwds == other._kwds
+        )
+
+    def __getattr__(self, attr: str):
+        # reminder: getattr is called after the normal lookups
+        if attr == "_kwds":
+            raise AttributeError()
+        if attr in self._kwds:
+            return self._kwds[attr]
+        else:
+            return getattr(super(), "__getattr__", super().__getattribute__)(attr)
+
+    def __getitem__(self, item: str) -> Any:
+        return self._kwds[item]
+
+    def __setattr__(self, item: str, val: Any) -> None:
+        if item.startswith("_"):
+            # Setting an instances copy of a ClassVar modify that
+            # By default, this makes **another** copy and places in _kwds
+            object.__setattr__(self, item, val)
+        else:
+            self._kwds[item] = val
+
+    def __setitem__(self, item: str, val: Any) -> None:
+        self._kwds[item] = val
+
+    def __repr__(self) -> str:
+        name = type(self).__name__
+        if kwds := self._kwds:
+            it = (f"{k}: {v!r}" for k, v in sorted(kwds.items()) if v is not Undefined)
+            args = ",\n".join(it).replace("\n", "\n  ")
+            LB, RB = "{", "}"
+            return f"{name}({LB}\n  {args}\n{RB})"
+        else:
+            return f"{name}({self._args[0]!r})"
+
+
 class SchemaBase:
     """
     Base class for schema wrappers.

tests/utils/test_schemapi.py

@@ -10,6 +10,7 @@
 import warnings
 from collections import deque
 from functools import partial
+from importlib.metadata import version as importlib_version
 from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence
 
 import jsonschema
@@ -18,6 +19,7 @@
 import pandas as pd
 import polars as pl
 import pytest
+from packaging.version import Version
 
 import altair as alt
 from altair import load_schema
@@ -51,6 +53,162 @@ def test_actual_json_schema_draft_is_same_as_hardcoded_default():
     )
 
 
+@pytest.fixture
+def dummy_rootschema() -> dict[str, Any]:
+    return {
+        "$schema": _JSON_SCHEMA_DRAFT_URL,
+        "definitions": {
+            "StringMapping": {
+                "type": "object",
+                "additionalProperties": {"type": "string"},
+            },
+            "StringArray": {"type": "array", "items": {"type": "string"}},
+        },
+        "properties": {
+            "a": {"$ref": "#/definitions/StringMapping"},
+            "a2": {"type": "object", "additionalProperties": {"type": "number"}},
+            "b": {"$ref": "#/definitions/StringArray"},
+            "b2": {"type": "array", "items": {"type": "number"}},
+            "c": {"type": ["string", "number"]},
+            "d": {
+                "anyOf": [
+                    {"$ref": "#/definitions/StringMapping"},
+                    {"$ref": "#/definitions/StringArray"},
+                ]
+            },
+            "e": {"items": [{"type": "string"}, {"type": "string"}]},
+        },
+    }
+
+
+def test_init_subclasses_hierarchy(dummy_rootschema) -> None:
+    from typing import Iterator
+
+    from altair.utils import schemapi
+
+    if Version(importlib_version("jsonschema")) >= Version("4.18"):
+        from referencing.exceptions import Unresolvable
+    else:
+        from jsonschema.exceptions import (  # type: ignore[assignment]
+            RefResolutionError as Unresolvable,
+        )
+
+    from altair.expr.core import GetItemExpression, OperatorMixin
+    from altair.utils.schemapi import _SchemaBasePEP487
+
+    sch1 = _SchemaBasePEP487()
+    sch2 = _SchemaBasePEP487()
+    sch3 = _SchemaBasePEP487("blue")
+    sch4 = _SchemaBasePEP487("red")
+    sch5 = _SchemaBasePEP487(color="blue")
+    sch6 = _SchemaBasePEP487(color="red")
+
+    with pytest.raises(
+        AssertionError, match=r"_SchemaBasePEP487\('blue', color='red'\)"
+    ):
+        _SchemaBasePEP487("blue", color="red")
+
+    assert sch1 == sch2
+    assert sch3 != sch4
+    assert sch5 != sch6
+    assert sch3 != sch5
+    assert _SchemaBasePEP487("blue") == sch3
+    assert _SchemaBasePEP487(color="red") == sch6
+    with pytest.raises(AttributeError, match="_SchemaBasePEP487.+color"):
+        attempt = sch4.color is Undefined  # noqa: F841
+
+    assert sch5.color == sch5["color"] == sch5._get("color") == "blue"
+    assert sch5._get("price") is Undefined
+    assert sch5._get("price", 999) == 999
+
+    assert _SchemaBasePEP487._class_is_valid_at_instantiation
+    sch6._class_is_valid_at_instantiation = False  # type: ignore[misc]
+    assert (
+        _SchemaBasePEP487._class_is_valid_at_instantiation
+        != sch6._class_is_valid_at_instantiation
+    )
+
+    with pytest.raises(TypeError, match="Test1PEP487.+ _schema"):
+
+        class Test1PEP487(_SchemaBasePEP487): ...
+
+    class Test2PEP487(_SchemaBasePEP487, schema={"type": "object"}): ...
+
+    with pytest.raises(
+        TypeError,
+        match=r"`rootschema` must be provided if `schema` contains a `'\$ref'` and does not inherit one",
+    ):
+
+        class Test3PEP487(_SchemaBasePEP487, schema={"$ref": "#/definitions/Bar"}): ...
+
+    class RootParentPEP487(_SchemaBasePEP487, schema=dummy_rootschema):
+        @classmethod
+        def _default_wrapper_classes(cls) -> Iterator[type[Any]]:
+            return schemapi._subclasses(RootParentPEP487)
+
+    class Root(RootParentPEP487):
+        """
+        Root schema wrapper.
+
+        A Vega-Lite top-level specification. This is the root class for all Vega-Lite
+        specifications. (The json schema is generated from this type.)
+        """
+
+        def __init__(self, *args, **kwds) -> None:
+            super().__init__(*args, **kwds)
+
+    assert (
+        Root._schema
+        == Root._rootschema
+        == RootParentPEP487._schema
+        == RootParentPEP487._rootschema
+    )
+
+    class StringMapping(Root, schema={"$ref": "#/definitions/StringMapping"}): ...
+
+    class StringArray(Root, schema={"$ref": "#/definitions/StringArray"}): ...
+
+    with pytest.raises(
+        jsonschema.ValidationError,
+        match=r"5 is not of type 'string'",
+    ):
+        schemapi.validate_jsonschema(
+            ["one", "two", 5], StringArray._schema, StringArray._rootschema
+        )
+
+    with pytest.raises(Unresolvable):
+        schemapi.validate_jsonschema(["one", "two", "three"], StringArray._schema)
+
+    schemapi.validate_jsonschema(
+        ["one", "two", "three"], StringArray._schema, StringArray._rootschema
+    )
+
+    class Expression(OperatorMixin, _SchemaBasePEP487, schema={"type": "string"}):
+        def to_dict(self, *args, **kwargs):
+            return repr(self)
+
+        def __setattr__(self, attr, val) -> None:
+            # We don't need the setattr magic defined in SchemaBase
+            return object.__setattr__(self, attr, val)
+
+        def __getitem__(self, val):
+            return GetItemExpression(self, val)
+
+    non_ref_mixin = Expression(
+        Expression("some").to_dict() + Expression("more").to_dict()
+    )
+    schemapi.validate_jsonschema(
+        non_ref_mixin.to_dict(), non_ref_mixin._schema, non_ref_mixin._rootschema
+    )
+    with pytest.raises(
+        jsonschema.ValidationError,
+        match=r"is not of type 'array'",
+    ):
+        schemapi.validate_jsonschema(
+            non_ref_mixin.to_dict(), StringArray._schema, StringArray._rootschema
+        )
+
+
 class _TestSchema(SchemaBase):
     @classmethod
     def _default_wrapper_classes(cls):