feat(typing): Adds public altair.typing module

altair/__init__.py

@@ -634,6 +634,7 @@
     "to_json",
     "to_values",
     "topo_feature",
+    "typing",
     "utils",
     "v5",
     "value",
@@ -654,6 +655,7 @@ def __dir__():
 from altair.jupyter import JupyterChart
 from altair.expr import expr
 from altair.utils import AltairDeprecationWarning, parse_shorthand, Optional, Undefined
+from altair import typing
 
 
 def load_ipython_extension(ipython):

altair/typing.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import sys
+from typing import Any, Mapping, Union
+from typing_extensions import TypedDict
+
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+ChannelType: TypeAlias = Union[str, Mapping[str, Any], Any]
+
+
+class EncodeKwds(TypedDict, total=False):
+    """
+    Reference implementation from [comment code block](https://github.com/pola-rs/polars/pull/17995#issuecomment-2263609625).
+
+    Aiming to define more specific `ChannelType`, without being exact.
+
+    This would be generated alongside `tools.generate_schema_wrapper._create_encode_signature`
+    """
+
+    angle: ChannelType
+    color: ChannelType
+    column: ChannelType
+    description: ChannelType
+    detail: ChannelType | list[Any]
+    facet: ChannelType
+    fill: ChannelType
+    fillOpacity: ChannelType
+    href: ChannelType
+    key: ChannelType
+    latitude: ChannelType
+    latitude2: ChannelType
+    longitude: ChannelType
+    longitude2: ChannelType
+    opacity: ChannelType
+    order: ChannelType | list[Any]
+    radius: ChannelType
+    radius2: ChannelType
+    row: ChannelType
+    shape: ChannelType
+    size: ChannelType
+    stroke: ChannelType
+    strokeDash: ChannelType
+    strokeOpacity: ChannelType
+    strokeWidth: ChannelType
+    text: ChannelType
+    theta: ChannelType
+    theta2: ChannelType
+    tooltip: ChannelType | list[Any]
+    url: ChannelType
+    x: ChannelType
+    x2: ChannelType
+    xError: ChannelType
+    xError2: ChannelType
+    xOffset: ChannelType
+    y: ChannelType
+    y2: ChannelType
+    yError: ChannelType
+    yError2: ChannelType
+    yOffset: ChannelType

altair/vegalite/v5/schema/channels.py

@@ -30905,7 +30905,7 @@ def encode(
         color: Optional[str | Color | Map | ColorDatum | ColorValue] = Undefined,
         column: Optional[str | Column | Map] = Undefined,
         description: Optional[str | Description | Map | DescriptionValue] = Undefined,
-        detail: Optional[str | Detail | Map | list] = Undefined,
+        detail: Optional[OneOrSeq[str | Detail | Map]] = Undefined,
         facet: Optional[str | Facet | Map] = Undefined,
         fill: Optional[str | Fill | Map | FillDatum | FillValue] = Undefined,
         fillOpacity: Optional[
@@ -30924,7 +30924,7 @@ def encode(
         opacity: Optional[
             str | Opacity | Map | OpacityDatum | OpacityValue
         ] = Undefined,
-        order: Optional[str | Order | Map | list | OrderValue] = Undefined,
+        order: Optional[OneOrSeq[str | Order | Map | OrderValue]] = Undefined,
         radius: Optional[str | Radius | Map | RadiusDatum | RadiusValue] = Undefined,
         radius2: Optional[
             str | Radius2 | Map | Radius2Datum | Radius2Value
@@ -30945,7 +30945,7 @@ def encode(
         text: Optional[str | Text | Map | TextDatum | TextValue] = Undefined,
         theta: Optional[str | Theta | Map | ThetaDatum | ThetaValue] = Undefined,
         theta2: Optional[str | Theta2 | Map | Theta2Datum | Theta2Value] = Undefined,
-        tooltip: Optional[str | Tooltip | Map | list | TooltipValue] = Undefined,
+        tooltip: Optional[OneOrSeq[str | Tooltip | Map | TooltipValue]] = Undefined,
         url: Optional[str | Url | Map | UrlValue] = Undefined,
         x: Optional[str | X | Map | XDatum | XValue] = Undefined,
         x2: Optional[str | X2 | Map | X2Datum | X2Value] = Undefined,

tools/generate_schema_wrapper.py

@@ -209,7 +209,7 @@ def configure_{prop}(self, *args, **kwargs) -> Self:
 
 ENCODE_METHOD: Final = '''
 class _EncodingMixin:
-    def encode({encode_method_args}) -> Self:
+    def encode({method_args}) -> Self:
         """Map properties of the data to visual properties of the chart (see :class:`FacetedEncoding`)
         {docstring}"""
         # Compat prep for `infer_encoding_types` signature
@@ -233,6 +233,13 @@ def encode({encode_method_args}) -> Self:
         return copy
 '''
 
+ENCODE_TYPED_DICT: Final = '''
+class _EncodeKwds(TypedDict, total=False):
+    """{docstring}"""
+    {channels}
+
+'''
+
 # NOTE: Not yet reasonable to generalize `TypeAliasType`, `TypeVar`
 # Revisit if this starts to become more common
 TYPING_EXTRA: Final = '''
@@ -562,18 +569,29 @@ def _type_checking_only_imports(*imports: str) -> str:
 class ChannelInfo:
     supports_arrays: bool
     deep_description: str
-    field_class_name: str | None = None
+    field_class_name: str
     datum_class_name: str | None = None
     value_class_name: str | None = None
 
+    @property
+    def is_field_only(self) -> bool:
+        return not (self.datum_class_name or self.value_class_name)
+
     @property
     def all_names(self) -> Iterator[str]:
-        if self.field_class_name:
-            yield self.field_class_name
-        if self.datum_class_name:
-            yield self.datum_class_name
-        if self.value_class_name:
-            yield self.value_class_name
+        """All channels are expected to have a field class."""
+        yield self.field_class_name
+        yield from self.non_field_names
+
+    @property
+    def non_field_names(self) -> Iterator[str]:
+        if self.is_field_only:
+            yield from ()
+        else:
+            if self.datum_class_name:
+                yield self.datum_class_name
+            if self.value_class_name:
+                yield self.value_class_name
 
 
 def generate_vegalite_channel_wrappers(
@@ -595,50 +613,37 @@ def generate_vegalite_channel_wrappers(
         supports_arrays = any(
             schema_info.is_array() for schema_info in propschema.anyOf
         )
+        classname: str = prop[0].upper() + prop[1:]
         channel_info = ChannelInfo(
             supports_arrays=supports_arrays,
             deep_description=propschema.deep_description,
+            field_class_name=classname,
         )
 
         for encoding_spec, definition in def_dict.items():
-            classname = prop[0].upper() + prop[1:]
             basename = definition.rsplit("/", maxsplit=1)[-1]
             basename = get_valid_identifier(basename)
 
+            gen: SchemaGenerator
             defschema = {"$ref": definition}
-
-            Generator: (
-                type[FieldSchemaGenerator]
-                | type[DatumSchemaGenerator]
-                | type[ValueSchemaGenerator]
-            )
+            kwds = {
+                "basename": basename,
+                "schema": defschema,
+                "rootschema": schema,
+                "encodingname": prop,
+                "haspropsetters": True,
+            }
             if encoding_spec == "field":
-                Generator = FieldSchemaGenerator
-                nodefault = []
-                channel_info.field_class_name = classname
-
+                gen = FieldSchemaGenerator(classname, nodefault=[], **kwds)
             elif encoding_spec == "datum":
-                Generator = DatumSchemaGenerator
-                classname += "Datum"
-                nodefault = ["datum"]
-                channel_info.datum_class_name = classname
-
+                temp_name = f"{classname}Datum"
+                channel_info.datum_class_name = temp_name
+                gen = DatumSchemaGenerator(temp_name, nodefault=["datum"], **kwds)
             elif encoding_spec == "value":
-                Generator = ValueSchemaGenerator
-                classname += "Value"
-                nodefault = ["value"]
-                channel_info.value_class_name = classname
-
-            gen = Generator(
-                classname=classname,
-                basename=basename,
-                schema=defschema,
-                rootschema=schema,
-                encodingname=prop,
-                nodefault=nodefault,
-                haspropsetters=True,
-                altair_classes_prefix="core",
-            )
+                temp_name = f"{classname}Value"
+                channel_info.value_class_name = temp_name
+                gen = ValueSchemaGenerator(temp_name, nodefault=["value"], **kwds)
+
             class_defs.append(gen.schema_class())
 
         channel_infos[prop] = channel_info
@@ -656,7 +661,7 @@ def generate_vegalite_channel_wrappers(
 
     imports = imports or [
         "from __future__ import annotations\n",
-        "from typing import Any, overload, Sequence, List, Literal, Union, TYPE_CHECKING",
+        "from typing import Any, overload, Sequence, List, Literal, Union, TYPE_CHECKING, TypedDict",
         "from narwhals.dependencies import is_pandas_dataframe as _is_pandas_dataframe",
         "from altair.utils.schemapi import Undefined, with_property_setters",
         "from altair.utils import infer_encoding_types as _infer_encoding_types",
@@ -676,11 +681,8 @@ def generate_vegalite_channel_wrappers(
         "\n" f"__all__ = {sorted(all_)}\n",
         CHANNEL_MIXINS,
         *class_defs,
+        *EncodingArtifacts(channel_infos, ENCODE_METHOD, ENCODE_TYPED_DICT),
     ]
-
-    # Generate the type signature for the encode method
-    encode_signature = _create_encode_signature(channel_infos)
-    contents.append(encode_signature)
     return "\n".join(contents)
 
 
@@ -861,59 +863,64 @@ def vegalite_main(skip_download: bool = False) -> None:
         ruff_write_lint_format_str(fp, contents)
 
 
-def _create_encode_signature(
-    channel_infos: dict[str, ChannelInfo],
-) -> str:
-    signature_args: list[str] = ["self", "*args: Any"]
-    docstring_parameters: list[str] = ["", "Parameters", "----------"]
-    for channel, info in channel_infos.items():
-        field_class_name = info.field_class_name
-        assert (
-            field_class_name is not None
-        ), "All channels are expected to have a field class"
-        datum_and_value_class_names = []
-        if info.datum_class_name is not None:
-            datum_and_value_class_names.append(info.datum_class_name)
-
-        if info.value_class_name is not None:
-            datum_and_value_class_names.append(info.value_class_name)
-
-        # dict stands for the return types of alt.datum, alt.value as well as
-        # the dictionary representation of an encoding channel class. See
-        # discussions in https://github.com/vega/altair/pull/3208
-        # for more background.
-        union_types = ["str", field_class_name, "Map"]
-        docstring_union_types = ["str", rst_syntax_for_class(field_class_name), "Dict"]
-        if info.supports_arrays:
-            # We could be more specific about what types are accepted in the list
-            # but then the signatures would get rather long and less useful
-            # to a user when it shows up in their IDE.
-            union_types.append("list")
-            docstring_union_types.append("List")
-
-        union_types = union_types + datum_and_value_class_names
-        docstring_union_types = docstring_union_types + [
-            rst_syntax_for_class(c) for c in datum_and_value_class_names
-        ]
+@dataclass
+class EncodingArtifacts:
+    """
+    Wrapper for what was previously `_create_encode_signature()`.
 
-        signature_args.append(
-            f"{channel}: Optional[Union[{', '.join(union_types)}]] = Undefined"
-        )
+    Now also creates a `TypedDict` as part of https://github.com/pola-rs/polars/pull/17995
+    """
 
-        docstring_parameters.extend(
-            (
-                f"{channel} : {', '.join(docstring_union_types)}",
-                f"    {process_description(info.deep_description)}",
+    channel_infos: dict[str, ChannelInfo]
+    fmt_method: str
+    fmt_typed_dict: str
+
+    def __iter__(self) -> Iterator[str]:
+        """After construction, this allows for unpacking (`*`)."""
+        yield from self._gen_artifacts()
+
+    def _gen_artifacts(self) -> None:
+        """
+        Generate `.encode()` related things.
+
+        Notes
+        -----
+        - `Map`/`Dict` stands for the return types of `alt.(datum|value)`, and any encoding channel class.
+            - See discussions in https://github.com/vega/altair/pull/3208
+        - We could be more specific about what types are accepted in the `List`
+            - but this translates poorly to an IDE
+            - `info.supports_arrays`
+        """
+        signature_args: list[str] = ["self", "*args: Any"]
+        docstring_parameters: list[str] = ["", "Parameters", "----------"]
+        typed_dict_args: list[str] = []
+        for channel, info in self.channel_infos.items():
+            it = info.all_names
+            it_rst_names = (rst_syntax_for_class(c) for c in info.all_names)
+
+            docstring_union_types = ["str", next(it_rst_names), "Dict"]
+            tp_inner: str = " | ".join(chain(("str", next(it), "Map"), it))
+            if info.supports_arrays:
+                docstring_union_types.append("List")
+                tp_inner = f"OneOrSeq[{tp_inner}]"
+            signature_args.append(f"{channel}: Optional[{tp_inner}] = Undefined")
+            typed_dict_args.append(f"{channel}: {tp_inner}")
+            docstring_parameters.extend(
+                (
+                    f"{channel} : {', '.join(chain(docstring_union_types, it_rst_names))}",
+                    f"    {process_description(info.deep_description)}",
+                )
             )
-        )
-    if len(docstring_parameters) > 1:
         docstring_parameters += [""]
-    docstring = indent_docstring(
-        docstring_parameters, indent_level=8, width=100, lstrip=False
-    )
-    return ENCODE_METHOD.format(
-        encode_method_args=", ".join(signature_args), docstring=docstring
-    )
+        doc = indent_docstring(
+            docstring_parameters, indent_level=8, width=100, lstrip=False
+        )
+        yield self.fmt_method.format(
+            method_args=", ".join(signature_args), docstring=doc
+        )
+        yield self.fmt_typed_dict.format(
+            channels="\n    ".join(typed_dict_args), docstring="Placeholder (FIXME)"
+        )
 
 
 def main() -> None: