diff --git a/altair/utils/_vegafusion_data.py b/altair/utils/_vegafusion_data.py index 66e3d99fb..21ca6833f 100644 --- a/altair/utils/_vegafusion_data.py +++ b/altair/utils/_vegafusion_data.py @@ -1,9 +1,13 @@ from __future__ import annotations import uuid +from importlib.metadata import version as importlib_version from typing import TYPE_CHECKING, Any, Callable, Final, TypedDict, Union, overload from weakref import WeakValueDictionary +from narwhals.dependencies import is_into_dataframe +from packaging.version import Version + from altair.utils._importers import import_vegafusion from altair.utils.core import DataFrameLike from altair.utils.data import ( @@ -15,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 @@ -33,6 +43,25 @@ VEGAFUSION_PREFIX: Final = "vegafusion+dataset://" +try: + VEGAFUSION_VERSION: Version | None = Version(importlib_version("vegafusion")) +except ImportError: + VEGAFUSION_VERSION = None + + +if VEGAFUSION_VERSION and Version("2.0.0a0") <= VEGAFUSION_VERSION: + + 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) + + class _ToVegaFusionReturnUrlDict(TypedDict): url: str @@ -64,7 +93,8 @@ def vegafusion_data_transformer( """VegaFusion Data Transformer.""" if data is None: return vegafusion_data_transformer - elif isinstance(data, DataFrameLike) 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] = data return {"url": VEGAFUSION_PREFIX + table_name}