diff --git a/altair/utils/_vegafusion_data.py b/altair/utils/_vegafusion_data.py index 8d39fb466..8b46bab78 100644 --- a/altair/utils/_vegafusion_data.py +++ b/altair/utils/_vegafusion_data.py @@ -2,13 +2,24 @@ import uuid from weakref import WeakValueDictionary -from typing import Union, Dict, Set, MutableMapping, TypedDict, Final, Any +from typing import ( + Union, + Dict, + Set, + MutableMapping, + TypedDict, + Final, + TYPE_CHECKING, +) from altair.utils._importers import import_vegafusion from altair.utils.core import DataFrameLike from altair.utils.data import DataType, ToValuesReturnType, MaxRowsError from altair.vegalite.data import default_data_transformer +if TYPE_CHECKING: + from vegafusion.runtime import ChartState # type: ignore + # Temporary storage for dataframes that have been extracted # from charts by the vegafusion data transformer. Use a WeakValueDictionary # rather than a dict so that the Python interpreter is free to garbage @@ -122,7 +133,9 @@ def get_inline_tables(vega_spec: dict) -> Dict[str, DataFrameLike]: return tables -def compile_to_vegafusion_chart_state(vegalite_spec: dict, local_tz: str) -> Any: +def compile_to_vegafusion_chart_state( + vegalite_spec: dict, local_tz: str +) -> "ChartState": """Compile a Vega-Lite spec to a VegaFusion ChartState Note: This function should only be called on a Vega-Lite spec @@ -223,7 +236,7 @@ def compile_with_vegafusion(vegalite_spec: dict) -> dict: return transformed_vega_spec -def handle_row_limit_exceeded(row_limit, warnings): +def handle_row_limit_exceeded(row_limit: int, warnings: list): for warning in warnings: if warning.get("type") == "RowLimitExceeded": raise MaxRowsError(