diff --git a/altair/utils/_vegafusion_data.py b/altair/utils/_vegafusion_data.py index 876cfc569..dc681155f 100644 --- a/altair/utils/_vegafusion_data.py +++ b/altair/utils/_vegafusion_data.py @@ -2,10 +2,10 @@ import uuid from importlib.metadata import version as importlib_version -from typing import TYPE_CHECKING, Any, Callable, Final, TypedDict, Union, cast, overload +from typing import TYPE_CHECKING, Any, Callable, Final, TypedDict, Union, overload from weakref import WeakValueDictionary -import narwhals.stable.v1 as nw +from narwhals.dependencies import is_into_dataframe from packaging.version import Version from altair.utils._importers import import_vegafusion @@ -19,12 +19,18 @@ from altair.vegalite.data import default_data_transformer if TYPE_CHECKING: + import sys from collections.abc import MutableMapping from narwhals.typing import IntoDataFrame from vegafusion.runtime import ChartState + if sys.version_info >= (3, 13): + from typing import TypeIs + else: + from typing_extensions import TypeIs + # 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 @@ -69,23 +75,9 @@ def vegafusion_data_transformer( if data is None: return vegafusion_data_transformer - # Test whether VegaFusion supports the data type - is_v2 = ( - Version(version) >= Version("2.0.0a0") - if (version := importlib_version("vegafusion")) - else False - ) - if is_v2: - # VegaFusion v2 support narwhals-compatible DataFrames - supported_by_vf = isinstance(data, DataFrameLike) or isinstance( - nw.from_native(data, pass_through=True), (nw.DataFrame, nw.LazyFrame) - ) - else: - supported_by_vf = isinstance(data, DataFrameLike) - - if supported_by_vf and not isinstance(data, SupportsGeoInterface): + if is_supported_by_vf(data) and not isinstance(data, SupportsGeoInterface): table_name = f"table_{uuid.uuid4()}".replace("-", "_") - extracted_inline_tables[table_name] = cast(DataFrameLike, data) + extracted_inline_tables[table_name] = data return {"url": VEGAFUSION_PREFIX + table_name} else: # Use default transformer for geo interface objects @@ -94,6 +86,21 @@ def vegafusion_data_transformer( return default_data_transformer(data) +if (version := importlib_version("vegafusion")) and Version(version) >= Version( + "2.0.0a0" +): + + def is_supported_by_vf(data: Any) -> TypeIs[DataFrameLike]: + # Test whether VegaFusion supports the data type + # VegaFusion v2 support narwhals-compatible DataFrames + return isinstance(data, DataFrameLike) or is_into_dataframe(data) + +else: + + def is_supported_by_vf(data: Any) -> TypeIs[DataFrameLike]: + return isinstance(data, DataFrameLike) + + def get_inline_table_names(vega_spec: dict[str, Any]) -> set[str]: """ Get a set of the inline datasets names in the provided Vega spec.