Type hints: Infer types from vegalite schema for autogenerated code

CONTRIBUTING.md

@@ -72,7 +72,7 @@ hatch run test
 ```
 
 
-This also runs the [`black`](https://black.readthedocs.io/) code formatter, [`ruff`](https://ruff.rs/) linter and [`mypy`](https://mypy-lang.org/) as type checker.
+This also runs the [`ruff`](https://ruff.rs/) linter and formatter as well as [`mypy`](https://mypy-lang.org/) as type checker.
 
 
 Study the output of any failed tests and try to fix the issues

altair/utils/schemapi.py

@@ -44,7 +44,7 @@
 else:
     from typing_extensions import Self
 
-_TSchemaBase = TypeVar("_TSchemaBase", bound="SchemaBase")
+_TSchemaBase = TypeVar("_TSchemaBase", bound=Type["SchemaBase"])
 
 ValidationErrorList = List[jsonschema.exceptions.ValidationError]
 GroupedValidationErrors = Dict[str, ValidationErrorList]

altair/vegalite/v5/api.py

@@ -18,6 +18,7 @@
 
 from .data import data_transformers
 from ... import utils, expr
+from ...expr import core as _expr_core
 from .display import renderers, VEGALITE_VERSION, VEGAEMBED_VERSION, VEGA_VERSION
 from .theme import themes
 from .compiler import vegalite_compilers
@@ -186,9 +187,12 @@ def _get_channels_mapping() -> TypingDict[TypingType[core.SchemaBase], str]:
 
 # -------------------------------------------------------------------------
 # Tools for working with parameters
-class Parameter(expr.core.OperatorMixin, object):
+class Parameter(_expr_core.OperatorMixin):
     """A Parameter object"""
 
+    # NOTE: If you change this class, make sure that the protocol in
+    # altair/vegalite/v5/schema/core.py is updated accordingly if needed.
+
     _counter: int = 0
 
     @classmethod
@@ -238,23 +242,23 @@ def __invert__(self):
         if self.param_type == "selection":
             return SelectionPredicateComposition({"not": {"param": self.name}})
         else:
-            return expr.core.OperatorMixin.__invert__(self)
+            return _expr_core.OperatorMixin.__invert__(self)
 
     def __and__(self, other):
         if self.param_type == "selection":
             if isinstance(other, Parameter):
                 other = {"param": other.name}
             return SelectionPredicateComposition({"and": [{"param": self.name}, other]})
         else:
-            return expr.core.OperatorMixin.__and__(self, other)
+            return _expr_core.OperatorMixin.__and__(self, other)
 
     def __or__(self, other):
         if self.param_type == "selection":
             if isinstance(other, Parameter):
                 other = {"param": other.name}
             return SelectionPredicateComposition({"or": [{"param": self.name}, other]})
         else:
-            return expr.core.OperatorMixin.__or__(self, other)
+            return _expr_core.OperatorMixin.__or__(self, other)
 
     def __repr__(self) -> str:
         return "Parameter({0!r}, {1})".format(self.name, self.param)
@@ -267,20 +271,20 @@ def _from_expr(self, expr) -> "ParameterExpression":
 
     def __getattr__(
         self, field_name: str
-    ) -> Union[expr.core.GetAttrExpression, "SelectionExpression"]:
+    ) -> Union[_expr_core.GetAttrExpression, "SelectionExpression"]:
         if field_name.startswith("__") and field_name.endswith("__"):
             raise AttributeError(field_name)
-        _attrexpr = expr.core.GetAttrExpression(self.name, field_name)
+        _attrexpr = _expr_core.GetAttrExpression(self.name, field_name)
         # If self is a SelectionParameter and field_name is in its
         # fields or encodings list, then we want to return an expression.
         if check_fields_and_encodings(self, field_name):
             return SelectionExpression(_attrexpr)
-        return expr.core.GetAttrExpression(self.name, field_name)
+        return _expr_core.GetAttrExpression(self.name, field_name)
 
     # TODO: Are there any special cases to consider for __getitem__?
     # This was copied from v4.
-    def __getitem__(self, field_name: str) -> expr.core.GetItemExpression:
-        return expr.core.GetItemExpression(self.name, field_name)
+    def __getitem__(self, field_name: str) -> _expr_core.GetItemExpression:
+        return _expr_core.GetItemExpression(self.name, field_name)
 
 
 # Enables use of ~, &, | with compositions of selection objects.
@@ -295,7 +299,7 @@ def __or__(self, other):
         return SelectionPredicateComposition({"or": [self.to_dict(), other.to_dict()]})
 
 
-class ParameterExpression(expr.core.OperatorMixin, object):
+class ParameterExpression(_expr_core.OperatorMixin, object):
     def __init__(self, expr) -> None:
         self.expr = expr
 
@@ -309,7 +313,7 @@ def _from_expr(self, expr) -> "ParameterExpression":
         return ParameterExpression(expr=expr)
 
 
-class SelectionExpression(expr.core.OperatorMixin, object):
+class SelectionExpression(_expr_core.OperatorMixin, object):
     def __init__(self, expr) -> None:
         self.expr = expr
 
@@ -346,9 +350,9 @@ def value(value, **kwargs) -> dict:
 def param(
     name: Optional[str] = None,
     value: Union[Any, UndefinedType] = Undefined,
-    bind: Union[core.Binding, str, UndefinedType] = Undefined,
+    bind: Union[core.Binding, UndefinedType] = Undefined,
     empty: Union[bool, UndefinedType] = Undefined,
-    expr: Union[str, core.Expr, expr.core.Expression, UndefinedType] = Undefined,
+    expr: Union[str, core.Expr, _expr_core.Expression, UndefinedType] = Undefined,
     **kwds,
 ) -> Parameter:
     """Create a named parameter.
@@ -365,7 +369,7 @@ def param(
     value : any (optional)
         The default value of the parameter. If not specified, the parameter
         will be created without a default value.
-    bind : :class:`Binding`, str (optional)
+    bind : :class:`Binding` (optional)
         Binds the parameter to an external input element such as a slider,
         selection list or radio button group.
     empty : boolean (optional)
@@ -421,9 +425,14 @@ def param(
             # If both 'value' and 'init' are set, we ignore 'init'.
             kwds.pop("init")
 
+    # ignore[arg-type] comment is needed because we can also pass _expr_core.Expression
     if "select" not in kwds:
         parameter.param = core.VariableParameter(
-            name=parameter.name, bind=bind, value=value, expr=expr, **kwds
+            name=parameter.name,
+            bind=bind,
+            value=value,
+            expr=expr,  # type: ignore[arg-type]
+            **kwds,
         )
         parameter.param_type = "variable"
     elif "views" in kwds:
@@ -503,7 +512,7 @@ def selection_interval(
     value: Union[Any, UndefinedType] = Undefined,
     bind: Union[core.Binding, str, UndefinedType] = Undefined,
     empty: Union[bool, UndefinedType] = Undefined,
-    expr: Union[str, core.Expr, expr.core.Expression, UndefinedType] = Undefined,
+    expr: Union[str, core.Expr, _expr_core.Expression, UndefinedType] = Undefined,
     encodings: Union[List[str], UndefinedType] = Undefined,
     on: Union[str, UndefinedType] = Undefined,
     clear: Union[str, bool, UndefinedType] = Undefined,
@@ -807,7 +816,7 @@ def condition(
     test_predicates = (str, expr.Expression, core.PredicateComposition)
 
     condition: TypingDict[
-        str, Union[bool, str, expr.core.Expression, core.PredicateComposition]
+        str, Union[bool, str, _expr_core.Expression, core.PredicateComposition]
     ]
     if isinstance(predicate, Parameter):
         if (
@@ -1376,7 +1385,9 @@ def project(
             spacing=spacing,
             tilt=tilt,
             translate=translate,
-            type=type,
+            # Ignore as we type here `type` as a str but in core.Projection
+            # it's a Literal with all options
+            type=type,  # type: ignore[arg-type]
             **kwds,
         )
         return self.properties(projection=projection)
@@ -1537,9 +1548,9 @@ def transform_calculate(
         self,
         as_: Union[str, core.FieldName, UndefinedType] = Undefined,
         calculate: Union[
-            str, core.Expr, expr.core.Expression, UndefinedType
+            str, core.Expr, _expr_core.Expression, UndefinedType
         ] = Undefined,
-        **kwargs: Union[str, core.Expr, expr.core.Expression],
+        **kwargs: Union[str, core.Expr, _expr_core.Expression],
     ) -> Self:
         """
         Add a :class:`CalculateTransform` to the schema.
@@ -1600,10 +1611,10 @@ def transform_calculate(
             )
         if as_ is not Undefined or calculate is not Undefined:
             dct = {"as": as_, "calculate": calculate}
-            self = self._add_transform(core.CalculateTransform(**dct))
+            self = self._add_transform(core.CalculateTransform(**dct))  # type: ignore[arg-type]
         for as_, calculate in kwargs.items():
             dct = {"as": as_, "calculate": calculate}
-            self = self._add_transform(core.CalculateTransform(**dct))
+            self = self._add_transform(core.CalculateTransform(**dct))  # type: ignore[arg-type]
         return self
 
     def transform_density(
@@ -1832,7 +1843,7 @@ def transform_filter(
         filter: Union[
             str,
             core.Expr,
-            expr.core.Expression,
+            _expr_core.Expression,
             core.Predicate,
             Parameter,
             core.PredicateComposition,
@@ -1866,7 +1877,7 @@ def transform_filter(
             elif isinstance(filter.empty, bool):
                 new_filter["empty"] = filter.empty
             filter = new_filter  # type: ignore[assignment]
-        return self._add_transform(core.FilterTransform(filter=filter, **kwargs))
+        return self._add_transform(core.FilterTransform(filter=filter, **kwargs))  # type: ignore[arg-type]
 
     def transform_flatten(
         self,
@@ -2068,7 +2079,13 @@ def transform_pivot(
         """
         return self._add_transform(
             core.PivotTransform(
-                pivot=pivot, value=value, groupby=groupby, limit=limit, op=op
+                # Ignore as we type here `op` as a str but in core.PivotTransform
+                # it's a Literal with all options
+                pivot=pivot,
+                value=value,
+                groupby=groupby,
+                limit=limit,
+                op=op,  # type: ignore[arg-type]
             )
         )
 
@@ -2318,7 +2335,7 @@ def transform_timeunit(
                 )
         if as_ is not Undefined:
             dct = {"as": as_, "timeUnit": timeUnit, "field": field}
-            self = self._add_transform(core.TimeUnitTransform(**dct))
+            self = self._add_transform(core.TimeUnitTransform(**dct))  # type: ignore[arg-type]
         for as_, shorthand in kwargs.items():
             dct = utils.parse_shorthand(
                 shorthand,
@@ -2330,7 +2347,7 @@ def transform_timeunit(
             dct["as"] = as_
             if "timeUnit" not in dct:
                 raise ValueError("'{}' must include a valid timeUnit".format(shorthand))
-            self = self._add_transform(core.TimeUnitTransform(**dct))
+            self = self._add_transform(core.TimeUnitTransform(**dct))  # type: ignore[arg-type]
         return self
 
     def transform_window(
@@ -2426,7 +2443,8 @@ def transform_window(
                     )
                 )
                 assert not isinstance(window, UndefinedType)  # For mypy
-                window.append(core.WindowFieldDef(**kwds))
+                # Ignore as core.WindowFieldDef has a Literal type hint with all options
+                window.append(core.WindowFieldDef(**kwds))  # type: ignore[arg-type]
 
         return self._add_transform(
             core.WindowTransform(
@@ -2607,7 +2625,7 @@ def resolve_scale(self, *args, **kwargs) -> Self:
 
 
 class _EncodingMixin:
-    @utils.use_signature(core.FacetedEncoding)
+    @utils.use_signature(channels._encode_signature)
     def encode(self, *args, **kwargs) -> Self:
         # Convert args to kwargs based on their types.
         kwargs = utils.infer_encoding_types(args, kwargs, channels)
@@ -2760,7 +2778,10 @@ def __init__(
         **kwargs,
     ) -> None:
         super(Chart, self).__init__(
-            data=data,
+            # Data type hints won't match with what TopLevelUnitSpec expects
+            # as there is some data processing happening when converting to
+            # a VL spec
+            data=data,  # type: ignore[arg-type]
             encoding=encoding,
             mark=mark,
             width=width,

altair/vegalite/v5/schema/__init__.py

@@ -1,5 +1,8 @@
 # ruff: noqa
+
 from .core import *
-from .channels import *
-SCHEMA_VERSION = 'v5.15.1'
-SCHEMA_URL = 'https://vega.github.io/schema/vega-lite/v5.15.1.json'
+from .channels import *  # type: ignore[assignment]
+
+SCHEMA_VERSION = "v5.15.1"
+
+SCHEMA_URL = "https://vega.github.io/schema/vega-lite/v5.15.1.json"