Optimized notebook dependency checks

Signed-off-by: pdmurray <peynmurray@gmail.com>

python/ray/_private/worker.py

@@ -96,7 +96,7 @@
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
 from ray.util.tracing.tracing_helper import _import_from_string
 from ray.widgets import Template
-from ray.widgets.util import ensure_ipywidgets_dep
+from ray.widgets.util import repr_with_fallback
 
 SCRIPT_MODE = 0
 WORKER_MODE = 1
@@ -1061,7 +1061,7 @@ def _repr_html_(self):
             context_table=self._context_table_template(),
         )
 
-    @ensure_ipywidgets_dep("8")
+    @repr_with_fallback(["ipywidgets", "8"])
     def _get_widget_bundle(self, **kwargs) -> Dict[str, Any]:
         """Get the mimebundle for the widget representation of the context.
 

python/ray/tests/test_widgets.py

@@ -2,137 +2,179 @@
 from unittest import mock
 
 import pytest
-from ray.widgets.util import repr_with_fallback
+from ray.widgets.util import repr_with_fallback, _can_display_ipywidgets
 
 
+@pytest.fixture
+def fancy_mimebundle():
+    return {
+        "fancy/mimetype": "A fancy repr",
+        "text/plain": "A simple repr",
+    }
+
+
+@mock.patch("ray.widgets.util._get_ipython_shell_name")
+@mock.patch("importlib.util.find_spec")
 @mock.patch("importlib.import_module")
-@mock.patch("ray.widgets.util.in_notebook")
-def test_ensure_notebook_dep_missing(
-    mock_in_notebook, mock_import_module, propagate_logs, caplog
+def test_repr_with_fallback_missing(
+    mock_import_module,
+    mock_find_spec,
+    mock_get_ipython_shell_name,
+    propagate_logs,
+    caplog,
+    fancy_mimebundle,
 ):
-    """Test that missing notebook dependencies trigger a warning."""
+    """Test that missing notebook dependencies trigger a log message."""
     caplog.set_level(logging.INFO)
 
-    class MockDep:
-        __version__ = "8.0.0"
-
     def raise_import_error(*args):
         raise ImportError
 
-    mock_import_module.return_value = MockDep()
     mock_import_module.side_effect = raise_import_error
-
-    mock_in_notebook.return_value = True
+    mock_find_spec.return_value = None
+    mock_get_ipython_shell_name.return_value = "ZMQInteractiveShell"
 
     class DummyObject:
-        @repr_with_fallback(["somedep", "8"])
-        def dummy_ipython_display(self):
-            return
+        def __repr__(self):
+            return "dummy repr"
 
-    DummyObject().dummy_ipython_display()
+        @repr_with_fallback(["somedep", "8"])
+        def _repr_mimebundle_(self, **kwargs):
+            return fancy_mimebundle
 
+    result = DummyObject()._repr_mimebundle_()
+    assert result == {"text/plain": "dummy repr"}
     assert "Missing packages:" in caplog.records[-1].msg
 
 
+@mock.patch("ray.widgets.util._get_ipython_shell_name")
+@mock.patch("importlib.util.find_spec")
 @mock.patch("importlib.import_module")
-@mock.patch("ray.widgets.util.in_notebook")
-def test_ensure_notebook_dep_outdated(
-    mock_in_notebook, mock_import_module, propagate_logs, caplog
+def test_repr_with_fallback_outdated(
+    mock_import_module,
+    mock_find_spec,
+    mock_get_ipython_shell_name,
+    propagate_logs,
+    caplog,
+    fancy_mimebundle,
 ):
-    """Test that outdated notebook dependencies trigger a warning."""
+    """Test that outdated notebook dependencies trigger a log message."""
     caplog.set_level(logging.INFO)
 
     class MockDep:
         __version__ = "7.0.0"
 
     mock_import_module.return_value = MockDep()
-
-    mock_in_notebook.return_value = True
+    mock_find_spec.return_value = "a valid import spec"
+    mock_get_ipython_shell_name.return_value = "ZMQInteractiveShell"
 
     class DummyObject:
-        @repr_with_fallback(["somedep", "8"])
-        def dummy_ipython_display():
-            return
+        def __repr__(self):
+            return "dummy repr"
 
-    DummyObject().dummy_ipython_display()
+        @repr_with_fallback(["somedep", "8"])
+        def _repr_mimebundle_(self, **kwargs):
+            return fancy_mimebundle
 
+    result = DummyObject()._repr_mimebundle_()
+    assert result == {"text/plain": "dummy repr"}
     assert "Outdated packages:" in caplog.records[-1].msg
 
 
+@mock.patch("ray.widgets.util._get_ipython_shell_name")
+@mock.patch("importlib.util.find_spec")
 @mock.patch("importlib.import_module")
-@mock.patch("ray.widgets.util.in_notebook")
-def test_ensure_notebook_valid(
-    mock_in_notebook, mock_import_module, propagate_logs, caplog
+def test_repr_with_fallback_valid(
+    mock_import_module,
+    mock_find_spec,
+    mock_get_ipython_shell_name,
+    propagate_logs,
+    caplog,
+    fancy_mimebundle,
 ):
-    """Test that valid notebook dependencies don't trigger a warning."""
+    """Test that valid notebook dependencies don't trigger a log message."""
     caplog.set_level(logging.INFO)
 
     class MockDep:
         __version__ = "8.0.0"
 
     mock_import_module.return_value = MockDep()
-
-    mock_in_notebook.return_value = True
+    mock_find_spec.return_value = "a valid import spec"
+    mock_get_ipython_shell_name.return_value = "ZMQInteractiveShell"
 
     class DummyObject:
+        def __repr__(self):
+            return "dummy repr"
+
         @repr_with_fallback(["somedep", "8"])
-        def dummy_ipython_display(self):
-            return
+        def _repr_mimebundle_(self, **kwargs):
+            return fancy_mimebundle
+
+    result = DummyObject()._repr_mimebundle_()
+    assert len(caplog.records) == 0
+    assert result == fancy_mimebundle
 
-    DummyObject().dummy_ipython_display()
 
+@mock.patch("ray.widgets.util._can_display_ipywidgets")
+@mock.patch("importlib.util.find_spec")
+@mock.patch("importlib.import_module")
+def test_repr_with_fallback_invalid_shell(
+    mock_import_module,
+    mock_find_spec,
+    mock_can_display_ipywidgets,
+    caplog,
+    fancy_mimebundle,
+):
+    """Test that the mimebundle is correctly stripped if run in an invalid shell."""
+    caplog.set_level(logging.INFO)
+
+    class MockDep:
+        __version__ = "8.0.0"
+
+    mock_import_module.return_value = MockDep()
+    mock_find_spec.return_value = "a valid import spec"
+    mock_can_display_ipywidgets.return_value = False
+
+    class DummyObject:
+        def __repr__(self):
+            return "dummy repr"
+
+        @repr_with_fallback(["somedep", "8"])
+        def _repr_mimebundle_(self, **kwargs):
+            return fancy_mimebundle
+
+    result = DummyObject()._repr_mimebundle_()
     assert len(caplog.records) == 0
+    assert result == {"text/plain": "dummy repr"}
 
 
+@mock.patch("ray.widgets.util._get_ipython_shell_name")
+@mock.patch("importlib.util.find_spec")
+@mock.patch("importlib.import_module")
 @pytest.mark.parametrize(
-    "kernel",
+    "shell,can_display",
     [
-        ("google.colab.kernel"),
-        ("normal.ipython.kernel"),
+        ("ZMQInteractiveShell", True),
+        ("google.colab.kernel", False),
+        ("TerminalInteractiveShell", False),
+        ("", False),
     ],
 )
-def test_repr_fallback_if_colab(kernel):
-    """Test that the mimebundle is correctly stripped if run in google colab."""
-    pytest.importorskip("IPython", reason="IPython is not installed.")
-    with mock.patch("IPython.get_ipython") as mock_get_ipython:
-        mock_get_ipython.return_value = kernel
-
-        class DummyObject:
-            @repr_with_fallback()
-            def _repr_mimebundle_(self, **kwargs):
-                return {
-                    "fancy/mimetype": "A fancy repr",
-                    "text/plain": "A simple repr",
-                }
-
-        obj = DummyObject()
-        result = obj._repr_mimebundle_()
-
-        assert "text/plain" in result
-        if "google.colab" in kernel:
-            assert len(result) == 1
-        else:
-            assert len(result) == 2
-            assert "fancy/mimetype" in result
-
-
-@mock.patch("ray.widgets.util.in_ipython_shell")
-def test_repr_fallback_if_ipython_shell(mock_in_ipython):
-    mock_in_ipython.return_value = True
-
-    class DummyObject:
-        @repr_with_fallback()
-        def _repr_mimebundle_(self, **kwargs):
-            return {
-                "fancy/mimetype": "A fancy repr",
-                "text/plain": "A simple repr",
-            }
+def test_can_display_ipywidgets(
+    mock_import_module,
+    mock_find_spec,
+    mock_get_ipython_shell_name,
+    shell,
+    can_display,
+):
+    class MockDep:
+        __version__ = "8.0.0"
 
-    obj = DummyObject()
-    result = obj._repr_mimebundle_()
+    mock_import_module.return_value = MockDep()
+    mock_find_spec.return_value = "a valid import spec"
+    mock_get_ipython_shell_name.return_value = shell
 
-    assert "text/plain" in result
-    assert len(result) == 1
+    assert _can_display_ipywidgets(["somedep", "8"], message="") == can_display
 
 
 if __name__ == "__main__":