fix(typing): Improve Then annotations, autocompletion, docs (#3567)

Author: dangotbanned

altair/utils/__init__.py

@@ -12,14 +12,15 @@
 from .deprecation import AltairDeprecationWarning, deprecated, deprecated_warn
 from .html import spec_to_html
 from .plugin_registry import PluginRegistry
-from .schemapi import Optional, SchemaBase, Undefined, is_undefined
+from .schemapi import Optional, SchemaBase, SchemaLike, Undefined, is_undefined
 
 __all__ = (
     "SHORTHAND_KEYS",
     "AltairDeprecationWarning",
     "Optional",
     "PluginRegistry",
     "SchemaBase",
+    "SchemaLike",
     "Undefined",
     "deprecated",
     "deprecated_warn",

altair/utils/core.py

@@ -28,7 +28,7 @@
 from narwhals.dependencies import get_polars, is_pandas_dataframe
 from narwhals.typing import IntoDataFrame
 
-from altair.utils.schemapi import SchemaBase, Undefined
+from altair.utils.schemapi import SchemaBase, SchemaLike, Undefined
 
 if sys.version_info >= (3, 12):
     from typing import Protocol, TypeAliasType, runtime_checkable
@@ -869,6 +869,8 @@ def _wrap_in_channel(self, obj: Any, encoding: str, /):
             obj = {"shorthand": obj}
         elif isinstance(obj, (list, tuple)):
             return [self._wrap_in_channel(el, encoding) for el in obj]
+        elif isinstance(obj, SchemaLike):
+            obj = obj.to_dict()
         if channel := self.name_to_channel.get(encoding):
             tp = channel["value" if "value" in obj else "field"]
             try:

altair/utils/schemapi.py

@@ -18,10 +18,12 @@
     Any,
     Dict,
     Final,
+    Generic,
     Iterable,
     Iterator,
     List,
     Literal,
+    Mapping,
     Sequence,
     TypeVar,
     Union,
@@ -41,6 +43,11 @@
 # not yet be fully instantiated in case your code is being executed during import time
 from altair import vegalite
 
+if sys.version_info >= (3, 12):
+    from typing import Protocol, TypeAliasType, runtime_checkable
+else:
+    from typing_extensions import Protocol, TypeAliasType, runtime_checkable
+
 if TYPE_CHECKING:
     from types import ModuleType
     from typing import ClassVar
@@ -524,11 +531,7 @@ def _todict(obj: Any, context: dict[str, Any] | None, np_opt: Any, pd_opt: Any)
             for k, v in obj.items()
             if v is not Undefined
         }
-    elif (
-        hasattr(obj, "to_dict")
-        and (module_name := obj.__module__)
-        and module_name.startswith("altair")
-    ):
+    elif isinstance(obj, SchemaLike):
         return obj.to_dict()
     elif pd_opt is not None and isinstance(obj, pd_opt.Timestamp):
         return pd_opt.Timestamp(obj).isoformat()
@@ -789,6 +792,95 @@ def _get_default_error_message(
         return message
 
 
+_JSON_VT_co = TypeVar(
+    "_JSON_VT_co",
+    Literal["string"],
+    Literal["object"],
+    Literal["array"],
+    covariant=True,
+)
+"""
+One of a subset of JSON Schema `primitive types`_:
+
+    ["string", "object", "array"]
+
+.. _primitive types:
+    https://json-schema.org/draft-07/json-schema-validation#rfc.section.6.1.1
+"""
+
+_TypeMap = TypeAliasType(
+    "_TypeMap", Mapping[Literal["type"], _JSON_VT_co], type_params=(_JSON_VT_co,)
+)
+"""
+A single item JSON Schema using the `type`_ keyword.
+
+This may represent **one of**:
+
+    {"type": "string"}
+    {"type": "object"}
+    {"type": "array"}
+
+.. _type:
+    https://json-schema.org/understanding-json-schema/reference/type
+"""
+
+# NOTE: Type checkers want opposing things:
+# - `mypy`   : Covariant type variable "_JSON_VT_co" used in protocol where invariant one is expected  [misc]
+# - `pyright`: Type variable "_JSON_VT_co" used in generic protocol "SchemaLike" should be covariant [reportInvalidTypeVarUse]
+# Siding with `pyright` as this is consistent with https://github.com/python/typeshed/blob/9e506eb5e8fc2823db8c60ad561b1145ff114947/stdlib/typing.pyi#L690
+
+
+@runtime_checkable
+class SchemaLike(Generic[_JSON_VT_co], Protocol):  # type: ignore[misc]
+    """
+    Represents ``altair`` classes which *may* not derive ``SchemaBase``.
+
+    Attributes
+    ----------
+    _schema
+        A single item JSON Schema using the `type`_ keyword.
+
+    Notes
+    -----
+    Should be kept tightly defined to the **minimum** requirements for:
+        - Converting into a form that can be validated by `jsonschema`_.
+        - Avoiding calling ``.to_dict()`` on a class external to ``altair``.
+    - ``_schema`` is more accurately described as a ``ClassVar``
+        - See `discussion`_ for blocking issue.
+
+    .. _jsonschema:
+        https://github.com/python-jsonschema/jsonschema
+    .. _type:
+        https://json-schema.org/understanding-json-schema/reference/type
+    .. _discussion:
+        https://github.com/python/typing/discussions/1424
+    """
+
+    _schema: _TypeMap[_JSON_VT_co]
+
+    def to_dict(self, *args, **kwds) -> Any: ...
+
+
+@runtime_checkable
+class ConditionLike(SchemaLike[Literal["object"]], Protocol):
+    """
+    Represents the wrapped state of a conditional encoding or property.
+
+    Attributes
+    ----------
+    condition
+        One or more (predicate, statement) pairs which each form a condition.
+
+    Notes
+    -----
+    - Can be extended with additional conditions.
+    - *Does not* define a default value, but can be finalized with one.
+    """
+
+    condition: Any
+    _schema: _TypeMap[Literal["object"]] = {"type": "object"}
+
+
 class UndefinedType:
     """A singleton object for marking undefined parameters."""