Support enabling/disabling VegaFusion renderers as context managers

python/vegafusion-jupyter/tests/test_conext_manager.py

@@ -0,0 +1,52 @@
+import altair as alt
+import vegafusion as vf
+
+
+def test_enabler_context_manager():
+    alt.data_transformers.enable("json")
+    alt.renderers.enable("mimetype")
+
+    assert alt.data_transformers.active == "json"
+    assert alt.renderers.active == "mimetype"
+
+    with vf.enable_widget():
+        assert alt.data_transformers.active == "vegafusion-feather"
+        assert alt.renderers.active == "vegafusion-widget"
+
+    assert alt.data_transformers.active == "json"
+    assert alt.renderers.active == "mimetype"
+
+    ctx = vf.enable_widget()
+    assert alt.data_transformers.active == "vegafusion-feather"
+    assert alt.renderers.active == "vegafusion-widget"
+    assert repr(ctx) == "vegafusion.enable_widget()"
+
+
+def test_enabler_context_manager_preserves_options():
+    # No options
+    with vf.enable_widget():
+        assert alt.data_transformers.active == "vegafusion-feather"
+        assert alt.data_transformers.options == {"data_dir": "_vegafusion_data"}
+
+        assert alt.renderers.active == "vegafusion-widget"
+        assert alt.renderers.options == {
+            'debounce_max_wait': 60, 'debounce_wait': 30, 'download_source_link': None
+        }
+
+        # Check that Widget uses default options
+        widget = vf.jupyter.VegaFusionWidget()
+        assert widget.debounce_max_wait == 60
+
+    # Options in the enable call
+    with vf.enable_widget(debounce_max_wait=200, data_dir="my_data_dir"):
+        assert alt.data_transformers.active == "vegafusion-feather"
+        assert alt.data_transformers.options == {"data_dir": "my_data_dir"}
+
+        assert alt.renderers.active == "vegafusion-widget"
+        assert alt.renderers.options == {
+            'debounce_max_wait': 200, 'debounce_wait': 30, 'download_source_link': None
+        }
+
+        # Check that widget picks up options from context manager
+        widget = vf.jupyter.VegaFusionWidget()
+        assert widget.debounce_max_wait == 200

python/vegafusion-jupyter/vegafusion_jupyter/__init__.py

@@ -6,6 +6,8 @@
 # this program the details of the active license.
 
 import altair as alt
+
+from vegafusion import RendererTransformerEnabler
 from . import renderer
 from . import widget
 from .widget import VegaFusionWidget
@@ -27,14 +29,17 @@ def enable(
     # Import vegafusion.transformer so that vegafusion-feather transform
     # will be registered
     import vegafusion.transformer
-    alt.renderers.enable(
-        'vegafusion-widget',
-        debounce_wait=debounce_wait,
-        debounce_max_wait=debounce_max_wait,
-        download_source_link=download_source_link,
-    )
-    alt.data_transformers.enable(
-        'vegafusion-feather', data_dir=data_dir
+    return RendererTransformerEnabler(
+        renderer_ctx=alt.renderers.enable(
+            'vegafusion-widget',
+            debounce_wait=debounce_wait,
+            debounce_max_wait=debounce_max_wait,
+            download_source_link=download_source_link,
+        ),
+        data_transformer_ctx=alt.data_transformers.enable(
+            'vegafusion-feather', data_dir=data_dir
+        ),
+        repr_str=f"vegafusion.enable_widget()"
     )
 
 

python/vegafusion-jupyter/vegafusion_jupyter/widget.py

@@ -42,6 +42,7 @@ class VegaFusionWidget(DOMWidget):
 
     def __init__(self, *args, **kwargs):
         # Make sure transformer and renderer and registered
+        import vegafusion as vf
         import vegafusion.transformer
         import vegafusion_jupyter.renderer
 
@@ -60,11 +61,11 @@ def __init__(self, *args, **kwargs):
             else:
                 data_transformer_opts = dict()
 
-            with alt.renderers.enable("vegafusion-widget"):
-                with alt.data_transformers.enable("vegafusion-feather", **data_transformer_opts):
-                    # Temporarily enable the vegafusion renderer and transformer so
-                    # that we use them even if they are not enabled globally
-                    spec = spec.to_dict()
+            with vf.enable_widget(**data_transformer_opts):
+                # Temporarily enable the vegafusion renderer and transformer so
+                # that we use them even if they are not enabled globally
+                spec = spec.to_dict()
+
             # Set spec as a dict, which will be converted to string below
             kwargs["spec"] = spec
 

python/vegafusion/tests/test_conext_manager.py

@@ -0,0 +1,34 @@
+import altair as alt
+import vegafusion as vf
+
+
+def test_enabler_context_manager():
+    alt.data_transformers.enable("json")
+    alt.renderers.enable("mimetype")
+
+    assert alt.data_transformers.active == "json"
+    assert alt.renderers.active == "mimetype"
+
+    with vf.disable():
+        assert alt.data_transformers.active == "default"
+        assert alt.renderers.active == "default"
+
+    assert alt.data_transformers.active == "json"
+    assert alt.renderers.active == "mimetype"
+
+    ctx = vf.disable()
+    assert alt.data_transformers.active == "default"
+    assert alt.renderers.active == "default"
+    assert repr(ctx) == "vegafusion.disable()"
+
+    with vf.enable_mime():
+        assert alt.data_transformers.active == "vegafusion-inline"
+        assert alt.renderers.active == "vegafusion-mime"
+
+    assert alt.data_transformers.active == "default"
+    assert alt.renderers.active == "default"
+
+    ctx = vf.enable_mime()
+    assert alt.data_transformers.active == "vegafusion-inline"
+    assert alt.renderers.active == "vegafusion-mime"
+    assert repr(ctx) == "vegafusion.enable_mime(mimetype='html', embed_options=None)"

python/vegafusion/vegafusion/__init__.py

@@ -13,6 +13,8 @@
 import altair as alt
 
 from ._version import __version__
+from .utils import RendererTransformerEnabler
+
 # Import optional subpackages
 try:
     import vegafusion.jupyter
@@ -59,10 +61,13 @@ def enable_mime(mimetype="html", embed_options=None):
         Dictionary of options to pass to the vega-embed. Default
         entry is {'mode': 'vega'}.
     """
-    # Import vegafusion.transformer so that vegafusion-inline transform
-    # will be registered
-    alt.renderers.enable('vegafusion-mime', mimetype=mimetype, embed_options=embed_options)
-    alt.data_transformers.enable('vegafusion-inline')
+    return RendererTransformerEnabler(
+        renderer_ctx=alt.renderers.enable(
+            'vegafusion-mime', mimetype=mimetype, embed_options=embed_options
+        ),
+        data_transformer_ctx=alt.data_transformers.enable('vegafusion-inline'),
+        repr_str=f"vegafusion.enable_mime(mimetype={repr(mimetype)}, embed_options={repr(embed_options)})"
+    )
 
 
 def enable_widget(
@@ -72,7 +77,7 @@ def enable_widget(
     data_dir="_vegafusion_data"
 ):
     from vegafusion.jupyter import enable
-    enable(
+    return enable(
         download_source_link=download_source_link,
         debounce_wait=debounce_wait,
         debounce_max_wait=debounce_max_wait,
@@ -95,5 +100,8 @@ def disable():
 
     This does not affect the behavior of VegaFusionWidget
     """
-    alt.renderers.enable('default')
-    alt.data_transformers.enable('default')
+    return RendererTransformerEnabler(
+        renderer_ctx=alt.renderers.enable('default'),
+        data_transformer_ctx=alt.data_transformers.enable('default'),
+        repr_str="vegafusion.disable()"
+    )

python/vegafusion/vegafusion/utils.py

@@ -0,0 +1,24 @@
+class RendererTransformerEnabler:
+    """
+    Context manager for enabling a renderer and data
+    transformer together
+    """
+    def __init__(self, renderer_ctx, data_transformer_ctx, repr_str=None):
+        self.renderer_ctx = renderer_ctx
+        self.data_transformer_ctx = data_transformer_ctx
+        self.repr_str = repr_str
+
+    def __enter__(self):
+        self.renderer_ctx.__enter__()
+        self.data_transformer_ctx.__enter__()
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        self.renderer_ctx.__exit__(*args, **kwargs)
+        self.data_transformer_ctx.__exit__(*args, **kwargs)
+
+    def __repr__(self):
+        if self.repr_str:
+            return self.repr_str
+        else:
+            return "vegafusion.utils.RendererTransformerEnabler"