fix(typing): Resolve LSP violation for ChartType.data (#3740)

vega · Jan 2, 2025 · b6253ff · b6253ff
1 parent 709bbdf
commit b6253ff
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 27 deletions.
diff --git a/altair/vegalite/v5/api.py b/altair/vegalite/v5/api.py
@@ -3939,11 +3939,8 @@ def __init__(
         height: Optional[int | dict | Step | Literal["container"]] = Undefined,
         **kwargs: Any,
     ) -> None:
-        # Data type hints won't match with what TopLevelUnitSpec expects
-        # as there is some data processing happening when converting to
-        # a VL spec
         super().__init__(
-            data=data,  # type: ignore[arg-type]
+            data=data,
             encoding=encoding,
             mark=mark,
             width=width,
@@ -4328,7 +4325,7 @@ def __init__(
     ) -> None:
         for spec in concat:
             _check_if_valid_subspec(spec, "ConcatChart")
-        super().__init__(data=data, concat=list(concat), columns=columns, **kwargs)  # type: ignore[arg-type]
+        super().__init__(data=data, concat=list(concat), columns=columns, **kwargs)
         self.concat: list[ChartType]
         self.params: Optional[Sequence[_Parameter]]
         self.data: Optional[ChartDataType]
@@ -4432,7 +4429,7 @@ def __init__(
     ) -> None:
         for spec in hconcat:
             _check_if_valid_subspec(spec, "HConcatChart")
-        super().__init__(data=data, hconcat=list(hconcat), **kwargs)  # type: ignore[arg-type]
+        super().__init__(data=data, hconcat=list(hconcat), **kwargs)
         self.hconcat: list[ChartType]
         self.params: Optional[Sequence[_Parameter]]
         self.data: Optional[ChartDataType]
@@ -4536,7 +4533,7 @@ def __init__(
     ) -> None:
         for spec in vconcat:
             _check_if_valid_subspec(spec, "VConcatChart")
-        super().__init__(data=data, vconcat=list(vconcat), **kwargs)  # type: ignore[arg-type]
+        super().__init__(data=data, vconcat=list(vconcat), **kwargs)
         self.vconcat: list[ChartType]
         self.params: Optional[Sequence[_Parameter]]
         self.data: Optional[ChartDataType]
@@ -4644,7 +4641,7 @@ def __init__(
         for spec in layer:
             _check_if_valid_subspec(spec, "LayerChart")
             _check_if_can_be_layered(spec)
-        super().__init__(data=data, layer=list(layer), **kwargs)  # type: ignore[arg-type]
+        super().__init__(data=data, layer=list(layer), **kwargs)
         self.layer: list[ChartType]
         self.params: Optional[Sequence[_Parameter]]
         self.data: Optional[ChartDataType]
@@ -4775,7 +4772,7 @@ def __init__(
         _spec_as_list = [spec]
         params, _spec_as_list = _combine_subchart_params(params, _spec_as_list)
         spec = _spec_as_list[0]
-        super().__init__(data=data, spec=spec, facet=facet, params=params, **kwargs)  # type: ignore[arg-type]
+        super().__init__(data=data, spec=spec, facet=facet, params=params, **kwargs)
         self.data: Optional[ChartDataType]
         self.spec: ChartType
         self.params: Optional[Sequence[_Parameter]]

diff --git a/altair/vegalite/v5/schema/core.py b/altair/vegalite/v5/schema/core.py
@@ -20,6 +20,7 @@
 
     from altair import Parameter
     from altair.typing import Optional
+    from altair.vegalite.v5.api import ChartDataType
 
     from ._typing import *  # noqa: F403
 
@@ -7887,7 +7888,7 @@ class GenericUnitSpecEncodingAnyMark(VegaLiteSchema):
     def __init__(
         self,
         mark: Optional[SchemaBase | Map | Mark_T | CompositeMark_T] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         description: Optional[str] = Undefined,
         encoding: Optional[SchemaBase | Map] = Undefined,
         name: Optional[str] = Undefined,
@@ -11062,7 +11063,7 @@ class LookupData(VegaLiteSchema):
 
     def __init__(
         self,
-        data: Optional[SchemaBase | Map] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map] = Undefined,
         key: Optional[str | SchemaBase] = Undefined,
         fields: Optional[Sequence[str | SchemaBase]] = Undefined,
         **kwds,
@@ -21055,7 +21056,7 @@ def __init__(
         bounds: Optional[Literal["full", "flush"]] = Undefined,
         center: Optional[bool | SchemaBase | Map] = Undefined,
         columns: Optional[float] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         description: Optional[str] = Undefined,
         name: Optional[str] = Undefined,
         resolve: Optional[SchemaBase | Map] = Undefined,
@@ -21182,7 +21183,7 @@ def __init__(
         bounds: Optional[Literal["full", "flush"]] = Undefined,
         center: Optional[bool | SchemaBase | Map] = Undefined,
         columns: Optional[float] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         description: Optional[str] = Undefined,
         name: Optional[str] = Undefined,
         resolve: Optional[SchemaBase | Map] = Undefined,
@@ -21340,7 +21341,7 @@ def __init__(
         align: Optional[SchemaBase | Map | LayoutAlign_T] = Undefined,
         bounds: Optional[Literal["full", "flush"]] = Undefined,
         center: Optional[bool | SchemaBase | Map] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         description: Optional[str] = Undefined,
         encoding: Optional[SchemaBase | Map] = Undefined,
         height: Optional[float | SchemaBase | Literal["container"] | Map] = Undefined,
@@ -21429,7 +21430,7 @@ def __init__(
         hconcat: Optional[Sequence[SchemaBase | Map]] = Undefined,
         bounds: Optional[Literal["full", "flush"]] = Undefined,
         center: Optional[bool] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         description: Optional[str] = Undefined,
         name: Optional[str] = Undefined,
         resolve: Optional[SchemaBase | Map] = Undefined,
@@ -21537,7 +21538,7 @@ class LayerSpec(Spec, NonNormalizedSpec):
     def __init__(
         self,
         layer: Optional[Sequence[SchemaBase | Map]] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         description: Optional[str] = Undefined,
         encoding: Optional[SchemaBase | Map] = Undefined,
         height: Optional[float | SchemaBase | Literal["container"] | Map] = Undefined,
@@ -21677,7 +21678,7 @@ def __init__(
         bounds: Optional[Literal["full", "flush"]] = Undefined,
         center: Optional[bool | SchemaBase | Map] = Undefined,
         columns: Optional[float] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         description: Optional[str] = Undefined,
         name: Optional[str] = Undefined,
         resolve: Optional[SchemaBase | Map] = Undefined,
@@ -21807,7 +21808,7 @@ def __init__(
         bounds: Optional[Literal["full", "flush"]] = Undefined,
         center: Optional[bool | SchemaBase | Map] = Undefined,
         columns: Optional[float] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         description: Optional[str] = Undefined,
         name: Optional[str] = Undefined,
         resolve: Optional[SchemaBase | Map] = Undefined,
@@ -24400,7 +24401,7 @@ def __init__(
         center: Optional[bool | SchemaBase | Map] = Undefined,
         columns: Optional[float] = Undefined,
         config: Optional[SchemaBase | Map] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         datasets: Optional[SchemaBase | Map] = Undefined,
         description: Optional[str] = Undefined,
         name: Optional[str] = Undefined,
@@ -24565,7 +24566,7 @@ class TopLevelFacetSpec(TopLevelSpec):
 
     def __init__(
         self,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         facet: Optional[SchemaBase | Map] = Undefined,
         spec: Optional[SchemaBase | Map] = Undefined,
         align: Optional[SchemaBase | Map | LayoutAlign_T] = Undefined,
@@ -24704,7 +24705,7 @@ def __init__(
         bounds: Optional[Literal["full", "flush"]] = Undefined,
         center: Optional[bool] = Undefined,
         config: Optional[SchemaBase | Map] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         datasets: Optional[SchemaBase | Map] = Undefined,
         description: Optional[str] = Undefined,
         name: Optional[str] = Undefined,
@@ -24860,7 +24861,7 @@ def __init__(
             str | Parameter | SchemaBase | Map | ColorName_T
         ] = Undefined,
         config: Optional[SchemaBase | Map] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         datasets: Optional[SchemaBase | Map] = Undefined,
         description: Optional[str] = Undefined,
         encoding: Optional[SchemaBase | Map] = Undefined,
@@ -25067,7 +25068,7 @@ class TopLevelUnitSpec(TopLevelSpec):
 
     def __init__(
         self,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         mark: Optional[SchemaBase | Map | Mark_T | CompositeMark_T] = Undefined,
         align: Optional[SchemaBase | Map | LayoutAlign_T] = Undefined,
         autosize: Optional[SchemaBase | Map | AutosizeType_T] = Undefined,
@@ -25212,7 +25213,7 @@ def __init__(
         bounds: Optional[Literal["full", "flush"]] = Undefined,
         center: Optional[bool] = Undefined,
         config: Optional[SchemaBase | Map] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         datasets: Optional[SchemaBase | Map] = Undefined,
         description: Optional[str] = Undefined,
         name: Optional[str] = Undefined,
@@ -26229,7 +26230,7 @@ class UnitSpec(VegaLiteSchema):
     def __init__(
         self,
         mark: Optional[SchemaBase | Map | Mark_T | CompositeMark_T] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         description: Optional[str] = Undefined,
         encoding: Optional[SchemaBase | Map] = Undefined,
         name: Optional[str] = Undefined,
@@ -26333,7 +26334,7 @@ class UnitSpecWithFrame(VegaLiteSchema):
     def __init__(
         self,
         mark: Optional[SchemaBase | Map | Mark_T | CompositeMark_T] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         description: Optional[str] = Undefined,
         encoding: Optional[SchemaBase | Map] = Undefined,
         height: Optional[float | SchemaBase | Literal["container"] | Map] = Undefined,
@@ -26460,7 +26461,7 @@ def __init__(
         vconcat: Optional[Sequence[SchemaBase | Map]] = Undefined,
         bounds: Optional[Literal["full", "flush"]] = Undefined,
         center: Optional[bool] = Undefined,
-        data: Optional[SchemaBase | Map | None] = Undefined,
+        data: Optional[SchemaBase | ChartDataType | Map | None] = Undefined,
         description: Optional[str] = Undefined,
         name: Optional[str] = Undefined,
         resolve: Optional[SchemaBase | Map] = Undefined,

diff --git a/tools/generate_schema_wrapper.py b/tools/generate_schema_wrapper.py
@@ -269,6 +269,7 @@ class {name}(TypedDict{metaclass_kwds}):{comment}
 BIN: Literal["Bin"] = "Bin"
 IMPUTE: Literal["Impute"] = "Impute"
 INTO_CONDITION: Literal["IntoCondition"] = "IntoCondition"
+CHART_DATA_TYPE: Literal["ChartDataType"] = "ChartDataType"
 
 # NOTE: `core.py` typing imports
 DATETIME: Literal["DateTime"] = "DateTime"
@@ -761,6 +762,7 @@ def generate_vegalite_schema_wrapper(fp: Path, /) -> ModuleDef[str]:
             "from datetime import date, datetime",
             "from altair import Parameter",
             "from altair.typing import Optional",
+            f"from altair.vegalite.v5.api import {CHART_DATA_TYPE}",
             "from ._typing import * # noqa: F403",
         ),
         "\n" f"__all__ = {all_}\n",

diff --git a/tools/schemapi/utils.py b/tools/schemapi/utils.py
@@ -613,6 +613,8 @@ def title_to_type_reprs(self, *, use_concrete: bool) -> set[str]:
                 tps.add("Parameter")
             if self.is_datetime():
                 tps.add("Temporal")
+            if self.is_top_level_spec_data():
+                tps.add("ChartDataType")
         elif self.is_value():
             value = self.properties["value"]
             t = value.to_type_repr(target="annotation", use_concrete=use_concrete)
@@ -969,6 +971,9 @@ def is_theme_config_target(self) -> bool:
     def is_datetime(self) -> bool:
         return self.refname == "DateTime"
 
+    def is_top_level_spec_data(self) -> bool:
+        return self.refname == "Data"
+
 
 class Grouped(Generic[T]):
     """