From a0c091eca897e1b9eaa6bb70d8ae8dc1deff9e60 Mon Sep 17 00:00:00 2001
From: Jon Mease <jonmmease@gmail.com>
Date: Fri, 15 Nov 2024 11:08:35 -0500
Subject: [PATCH] feat: VegaFusion 2 will support narwhals (#3682)

---
 altair/utils/_vegafusion_data.py | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

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}