Update VegaFusion version check for 2.0+

vega · Nov 12, 2024 · daf958f · daf958f
1 parent 255f736
commit daf958f
Showing 1 changed file with 15 additions and 12 deletions.
diff --git 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'