diff --git a/CHANGELOG.md b/CHANGELOG.md index 126ff280b4..23b6b561af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ This project adheres to [Semantic Versioning](http://semver.org/). ### Fixed - Fix issue where user-specified `color_continuous_scale` was ignored when template had `autocolorscale=True` [[#5439](https://github.com/plotly/plotly.py/pull/5439)], with thanks to @antonymilne for the contribution! +- Use presence of `COLAB_NOTEBOOK_ID` env var to enable Colab renderer instead of testing import of `google.colab` [[#5473](https://github.com/plotly/plotly.py/pull/5473)], with thanks to @kevineger for the contribution! - Update tests to be compatible with numpy 2.4 [[#5522](https://github.com/plotly/plotly.py/pull/5522)], with thanks to @thunze for the contribution! ### Updated diff --git a/plotly/io/_renderers.py b/plotly/io/_renderers.py index 4c21760bb0..8eea1fe8bf 100644 --- a/plotly/io/_renderers.py +++ b/plotly/io/_renderers.py @@ -488,13 +488,10 @@ def show(fig, renderer=None, validate=True, **kwargs): elif ipython and ipython.get_ipython(): # Try to detect environment so that we can enable a useful # default renderer - if not default_renderer: - try: - import google.colab # noqa: F401 - default_renderer = "colab" - except ImportError: - pass + # Check if we're running in a Colab web notebook + if not default_renderer and "COLAB_NOTEBOOK_ID" in os.environ: + default_renderer = "colab" # Check if we're running in a Kaggle notebook if not default_renderer and os.path.exists("/kaggle/input"): diff --git a/tests/test_io/test_renderers.py b/tests/test_io/test_renderers.py index ac204d5d9e..73ce23bd78 100644 --- a/tests/test_io/test_renderers.py +++ b/tests/test_io/test_renderers.py @@ -1,4 +1,5 @@ import json +import os import threading import time @@ -418,3 +419,75 @@ def test_missing_webbrowser_methods(fig1): finally: # restore everything after this test webbrowser.get = removed_webbrowser_get_method + + +def test_colab_renderer_when_env_var_is_set(): + """ + When COLAB_NOTEBOOK_ID is present the default renderer should be 'colab'. + """ + import importlib + import plotly.io._renderers as _renderers_mod + from plotly import optional_imports + + fake_ipython = MagicMock() + original_get_module = optional_imports.get_module + + def patched_get_module(name, *args, **kwargs): + if name in ("IPython", "IPython.display"): + return fake_ipython + return original_get_module(name, *args, **kwargs) + + original_default = pio.renderers.default + try: + with mock.patch.dict(os.environ, {"COLAB_NOTEBOOK_ID": "fake-id"}, clear=True): + with mock.patch.object( + optional_imports, "get_module", side_effect=patched_get_module + ): + importlib.reload(_renderers_mod) + assert _renderers_mod.renderers.default == "colab" + finally: + importlib.reload(_renderers_mod) + pio.renderers.default = original_default + + +def test_colab_renderer_when_env_var_is_absent(): + """ + Without COLAB_NOTEBOOK_ID the default renderer must not be 'colab', + even when ``google.colab`` is importable. + + Regression test for https://github.com/plotly/plotly.py/pull/5473. + """ + import sys + import types + import importlib + import plotly.io._renderers as _renderers_mod + from plotly import optional_imports + + fake_ipython = MagicMock() + original_get_module = optional_imports.get_module + + def patched_get_module(name, *args, **kwargs): + if name in ("IPython", "IPython.display"): + return fake_ipython + return original_get_module(name, *args, **kwargs) + + # Make google.colab importable so the old ``import google.colab`` + # approach would have chosen the colab renderer here. + fake_google = types.ModuleType("google") + fake_google_colab = types.ModuleType("google.colab") + + original_default = pio.renderers.default + try: + with mock.patch.dict(os.environ, {}, clear=True): + with mock.patch.dict( + sys.modules, + {"google": fake_google, "google.colab": fake_google_colab}, + ): + with mock.patch.object( + optional_imports, "get_module", side_effect=patched_get_module + ): + importlib.reload(_renderers_mod) + assert _renderers_mod.renderers.default != "colab" + finally: + importlib.reload(_renderers_mod) + pio.renderers.default = original_default