From b292ccfb88ac84cf716cc0116ef2047633ab2654 Mon Sep 17 00:00:00 2001
From: Jon Mease <jonmmease@gmail.com>
Date: Wed, 13 Nov 2024 07:53:31 -0500
Subject: [PATCH] chore: Prep for VegaFusion 2.0 (#3680)

---
 altair/utils/_importers.py | 27 +++++++++++++++------------
 pyproject.toml             |  2 +-
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/altair/utils/_importers.py b/altair/utils/_importers.py
index c02cc7011..efe48df8b 100644
--- a/altair/utils/_importers.py
+++ b/altair/utils/_importers.py
@@ -15,18 +15,21 @@ def import_vegafusion() -> ModuleType:
         import vegafusion as vf
 
         version = importlib_version("vegafusion")
-        embed_version = importlib_version("vegafusion-python-embed")
-        if version != embed_version or Version(version) < Version(min_version):
-            msg = (
-                "The versions of the vegafusion and vegafusion-python-embed packages must match\n"
-                f"and must be version {min_version} or greater.\n"
-                f"Found:\n"
-                f" - vegafusion=={version}\n"
-                f" - vegafusion-python-embed=={embed_version}\n"
-            )
-            raise RuntimeError(msg)
-
-        return vf
+        if Version(version) >= Version("2.0.0a0"):
+            # In VegaFusion 2.0 there is no vegafusion-python-embed package
+            return vf
+        else:
+            embed_version = importlib_version("vegafusion-python-embed")
+            if version != embed_version or Version(version) < Version(min_version):
+                msg = (
+                    "The versions of the vegafusion and vegafusion-python-embed packages must match\n"
+                    f"and must be version {min_version} or greater.\n"
+                    f"Found:\n"
+                    f" - vegafusion=={version}\n"
+                    f" - vegafusion-python-embed=={embed_version}\n"
+                )
+                raise RuntimeError(msg)
+            return vf
     except ImportError as err:
         msg = (
             'The "vegafusion" data transformer and chart.transformed_data feature requires\n'
diff --git a/pyproject.toml b/pyproject.toml
index 4132f0a25..bcac9e098 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -63,7 +63,7 @@ all = [
     "pandas>=1.1.3",
     "numpy",
     "pyarrow>=11",
-    "vegafusion[embed]>=1.6.6,<2",
+    "vegafusion[embed]>=1.6.6",
     "anywidget>=0.9.0",
     "altair_tiles>=0.3.0"
 ]