Merge branch 'main' into fix-pylance-imports

altair/utils/__init__.py

@@ -13,7 +13,7 @@
 from .html import spec_to_html
 from .plugin_registry import PluginRegistry
 from .deprecation import AltairDeprecationWarning, deprecated, deprecated_warn
-from .schemapi import Undefined, Optional
+from .schemapi import Undefined, Optional, is_undefined
 
 
 __all__ = (
@@ -28,6 +28,7 @@
     "display_traceback",
     "infer_encoding_types",
     "infer_vegalite_type_for_pandas",
+    "is_undefined",
     "parse_shorthand",
     "sanitize_narwhals_dataframe",
     "sanitize_pandas_dataframe",

altair/utils/schemapi.py

@@ -787,6 +787,19 @@ def func_2(
 """
 
 
+def is_undefined(obj: Any) -> TypeIs[UndefinedType]:
+    """Type-safe singleton check for `UndefinedType`.
+
+    Notes
+    -----
+    - Using `obj is Undefined` does not narrow from `UndefinedType` in a union.
+        - Due to the assumption that other `UndefinedType`'s could exist.
+    - Current [typing spec advises](https://typing.readthedocs.io/en/latest/spec/concepts.html#support-for-singleton-types-in-unions) using an `Enum`.
+        - Otherwise, requires an explicit guard to inform the type checker.
+    """
+    return obj is Undefined
+
+
 class SchemaBase:
     """Base class for schema wrappers.
 

altair/vegalite/v5/api.py

@@ -14,7 +14,6 @@
     Union,
     TYPE_CHECKING,
     TypeVar,
-    Sequence,
     Protocol,
 )
 from typing_extensions import TypeAlias
@@ -46,10 +45,6 @@
     from typing import TypedDict
 else:
     from typing_extensions import TypedDict
-if sys.version_info >= (3, 12):
-    from typing import TypeAliasType
-else:
-    from typing_extensions import TypeAliasType
 
 if TYPE_CHECKING:
     from ...utils.core import DataFrameLike
@@ -108,7 +103,6 @@
         TopLevelSelectionParameter,
         SelectionParameter,
         InlineDataset,
-        UndefinedType,
     )
     from altair.expr.core import (
         BinaryExpression,
@@ -127,6 +121,7 @@
         AggregateOp_T,
         MultiTimeUnit_T,
         SingleTimeUnit_T,
+        OneOrSeq,
     )
 
 __all__ = [
@@ -184,21 +179,6 @@
 
 ChartDataType: TypeAlias = Optional[Union[DataType, core.Data, str, core.Generator]]
 _TSchemaBase = TypeVar("_TSchemaBase", bound=core.SchemaBase)
-_T = TypeVar("_T")
-_OneOrSeq = TypeAliasType("_OneOrSeq", Union[_T, Sequence[_T]], type_params=(_T,))
-"""One of ``_T`` specified type(s), or a `Sequence` of such.
-
-Examples
---------
-The parameters ``short``, ``long`` accept the same range of types::
-
-    # ruff: noqa: UP006, UP007
-
-    def func(
-        short: _OneOrSeq[str | bool | float],
-        long: Union[str, bool, float, Sequence[Union[str, bool, float]],
-    ): ...
-"""
 
 
 # ------------------------------------------------------------------------
@@ -234,7 +214,7 @@ def _consolidate_data(data: Any, context: Any) -> Any:
     kwds = {}
 
     if isinstance(data, core.InlineData):
-        if _is_undefined(data.name) and not _is_undefined(data.values):
+        if utils.is_undefined(data.name) and not utils.is_undefined(data.values):
             if isinstance(data.values, core.InlineDataset):
                 values = data.to_dict()["values"]
             else:
@@ -245,7 +225,7 @@ def _consolidate_data(data: Any, context: Any) -> Any:
         values = data["values"]
         kwds = {k: v for k, v in data.items() if k != "values"}
 
-    if not _is_undefined(values):
+    if not utils.is_undefined(values):
         name = _dataset_name(values)
         data = core.NamedData(name=name, **kwds)
         context.setdefault("datasets", {})[name] = values
@@ -469,7 +449,7 @@ def _from_expr(self, expr) -> SelectionExpression:
 
 def check_fields_and_encodings(parameter: Parameter, field_name: str) -> bool:
     param = parameter.param
-    if _is_undefined(param) or isinstance(param, core.VariableParameter):
+    if utils.is_undefined(param) or isinstance(param, core.VariableParameter):
         return False
     for prop in ["fields", "encodings"]:
         try:
@@ -538,19 +518,6 @@ def _is_test_predicate(obj: Any) -> TypeIs[_TestPredicateType]:
     return isinstance(obj, (str, _expr_core.Expression, core.PredicateComposition))
 
 
-def _is_undefined(obj: Any) -> TypeIs[UndefinedType]:
-    """Type-safe singleton check for `UndefinedType`.
-
-    Notes
-    -----
-    - Using `obj is Undefined` does not narrow from `UndefinedType` in a union.
-        - Due to the assumption that other `UndefinedType`'s could exist.
-    - Current [typing spec advises](https://typing.readthedocs.io/en/latest/spec/concepts.html#support-for-singleton-types-in-unions) using an `Enum`.
-        - Otherwise, requires an explicit guard to inform the type checker.
-    """
-    return obj is Undefined
-
-
 def _get_predicate_expr(p: Parameter) -> Optional[str | SchemaBase]:
     # https://vega.github.io/vega-lite/docs/predicate.html
     return getattr(p.param, "expr", Undefined)
@@ -562,7 +529,7 @@ def _predicate_to_condition(
     condition: _ConditionType
     if isinstance(predicate, Parameter):
         predicate_expr = _get_predicate_expr(predicate)
-        if predicate.param_type == "selection" or _is_undefined(predicate_expr):
+        if predicate.param_type == "selection" or utils.is_undefined(predicate_expr):
             condition = {"param": predicate.name}
             if isinstance(empty, bool):
                 condition["empty"] = empty
@@ -638,7 +605,7 @@ class _ConditionExtra(TypedDict, closed=True, total=False):  # type: ignore[call
     param: Parameter | str
     test: _TestPredicateType
     value: Any
-    __extra_items__: _StatementType | _OneOrSeq[_LiteralValue]
+    __extra_items__: _StatementType | OneOrSeq[_LiteralValue]
 
 
 _Condition: TypeAlias = _ConditionExtra
@@ -749,7 +716,7 @@ def _parse_when(
     **constraints: _FieldEqualType,
 ) -> _ConditionType:
     composed: _PredicateType
-    if _is_undefined(predicate):
+    if utils.is_undefined(predicate):
         if more_predicates or constraints:
             composed = _parse_when_compose(more_predicates, constraints)
         else:
@@ -1159,7 +1126,7 @@ def param(
     empty_remap = {"none": False, "all": True}
     parameter = Parameter(name)
 
-    if not _is_undefined(empty):
+    if not utils.is_undefined(empty):
         if isinstance(empty, bool) and not isinstance(empty, str):
             parameter.empty = empty
         elif empty in empty_remap:
@@ -1677,7 +1644,7 @@ def to_dict(
 
         copy = _top_schema_base(self).copy(deep=False)
         original_data = getattr(copy, "data", Undefined)
-        if not _is_undefined(original_data):
+        if not utils.is_undefined(original_data):
             try:
                 data = _to_eager_narwhals_dataframe(original_data)
             except TypeError:

altair/vegalite/v5/schema/_typing.py

@@ -4,9 +4,9 @@
 
 from __future__ import annotations
 
-from typing import Any, Literal, Mapping
+from typing import Any, Literal, Mapping, Sequence, TypeVar, Union
 
-from typing_extensions import TypeAlias
+from typing_extensions import TypeAlias, TypeAliasType
 
 __all__ = [
     "AggregateOp_T",
@@ -32,6 +32,7 @@
     "Mark_T",
     "MultiTimeUnit_T",
     "NonArgAggregateOp_T",
+    "OneOrSeq",
     "Orient_T",
     "Orientation_T",
     "ProjectionType_T",
@@ -60,6 +61,22 @@
 ]
 
 
+T = TypeVar("T")
+OneOrSeq = TypeAliasType("OneOrSeq", Union[T, Sequence[T]], type_params=(T,))
+"""One of ``T`` specified type(s), or a `Sequence` of such.
+
+Examples
+--------
+The parameters ``short``, ``long`` accept the same range of types::
+
+    # ruff: noqa: UP006, UP007
+
+    def func(
+        short: OneOrSeq[str | bool | float],
+        long: Union[str, bool, float, Sequence[Union[str, bool, float]],
+    ): ...
+"""
+
 Map: TypeAlias = Mapping[str, Any]
 AggregateOp_T: TypeAlias = Literal[
     "argmax",

tools/generate_schema_wrapper.py

@@ -233,6 +233,26 @@ def encode({encode_method_args}) -> Self:
         return copy
 '''
 
+# NOTE: Not yet reasonable to generalize `TypeAliasType`, `TypeVar`
+# Revisit if this starts to become more common
+TYPING_EXTRA: Final = '''
+T = TypeVar("T")
+OneOrSeq = TypeAliasType("OneOrSeq", Union[T, Sequence[T]], type_params=(T,))
+"""One of ``T`` specified type(s), or a `Sequence` of such.
+
+Examples
+--------
+The parameters ``short``, ``long`` accept the same range of types::
+
+    # ruff: noqa: UP006, UP007
+
+    def func(
+        short: OneOrSeq[str | bool | float],
+        long: Union[str, bool, float, Sequence[Union[str, bool, float]],
+    ): ...
+"""
+'''
+
 
 class SchemaGenerator(codegen.SchemaGenerator):
     schema_class_template = textwrap.dedent(
@@ -839,7 +859,9 @@ def vegalite_main(skip_download: bool = False) -> None:
     )
     print(msg)
     TypeAliasTracer.update_aliases(("Map", "Mapping[str, Any]"))
-    TypeAliasTracer.write_module(fp_typing, header=HEADER)
+    TypeAliasTracer.write_module(
+        fp_typing, "OneOrSeq", header=HEADER, extra=TYPING_EXTRA
+    )
     # Write the pre-generated modules
     for fp, contents in files.items():
         print(f"Writing\n {schemafile!s}\n  ->{fp!s}")

tools/schemapi/schemapi.py

@@ -785,6 +785,19 @@ def func_2(
 """
 
 
+def is_undefined(obj: Any) -> TypeIs[UndefinedType]:
+    """Type-safe singleton check for `UndefinedType`.
+
+    Notes
+    -----
+    - Using `obj is Undefined` does not narrow from `UndefinedType` in a union.
+        - Due to the assumption that other `UndefinedType`'s could exist.
+    - Current [typing spec advises](https://typing.readthedocs.io/en/latest/spec/concepts.html#support-for-singleton-types-in-unions) using an `Enum`.
+        - Otherwise, requires an explicit guard to inform the type checker.
+    """
+    return obj is Undefined
+
+
 class SchemaBase:
     """Base class for schema wrappers.
 

tools/schemapi/utils.py

@@ -71,8 +71,8 @@ def __init__(
         self._aliases: dict[str, str] = {}
         self._imports: Sequence[str] = (
             "from __future__ import annotations\n",
-            "from typing import Literal, Mapping, Any",
-            "from typing_extensions import TypeAlias",
+            "from typing import Any, Literal, Mapping, TypeVar, Sequence, Union",
+            "from typing_extensions import TypeAlias, TypeAliasType",
         )
         self._cmd_check: list[str] = ["--fix"]
         self._cmd_format: Sequence[str] = ruff_format or ()
@@ -141,28 +141,31 @@ def is_cached(self, tp: str, /) -> bool:
         return tp in self._literals_invert or tp in self._literals
 
     def write_module(
-        self, fp: Path, *extra_imports: str, header: LiteralString
+        self, fp: Path, *extra_all: str, header: LiteralString, extra: LiteralString
     ) -> None:
         """Write all collected `TypeAlias`'s to `fp`.
 
         Parameters
         ----------
         fp
             Path to new module.
-        *extra_imports
-            Follows `self._imports` block.
+        *extra_all
+            Any manually spelled types to be exported.
         header
             `tools.generate_schema_wrapper.HEADER`.
+        extra
+            `tools.generate_schema_wrapper.TYPING_EXTRA`.
         """
         ruff_format = ["ruff", "format", fp]
         if self._cmd_format:
             ruff_format.extend(self._cmd_format)
         commands = (["ruff", "check", fp, *self._cmd_check], ruff_format)
-        static = (header, "\n", *self._imports, *extra_imports, "\n\n")
+        static = (header, "\n", *self._imports, "\n\n")
         self.update_aliases(*sorted(self._literals.items(), key=itemgetter(0)))
+        all_ = [*iter(self._aliases), *extra_all]
         it = chain(
             static,
-            [f"__all__ = {list(self._aliases)}", "\n\n"],
+            [f"__all__ = {all_}", "\n\n", extra],
             self.generate_aliases(),
         )
         fp.write_text("\n".join(it), encoding="utf-8")