simpler type hints when plotly is not installed

Author: luismavs

Makefile

@@ -4,7 +4,7 @@ test:
 test-no-plotly:
 	uv sync --extra test
 	uv pip uninstall plotly
-	pytest tests/test_other.py -k plotly
+	pytest tests/test_other.py -k plotly --pdb
 	uv sync --extra test
 	pytest tests/test_other.py -k plotly
 

bertopic/_bertopic.py

@@ -1,6 +1,7 @@
 # ruff: noqa: E402
 import yaml
 import warnings
+import importlib
 
 warnings.filterwarnings("ignore", category=FutureWarning)
 warnings.filterwarnings("ignore", category=UserWarning)
@@ -26,6 +27,7 @@
 from collections import defaultdict, Counter
 from scipy.sparse import csr_matrix
 from scipy.cluster import hierarchy as sch
+from importlib.util import find_spec
 
 # Typing
 import sys
@@ -34,7 +36,19 @@
     from typing import Literal
 else:
     from typing_extensions import Literal
-from typing import List, Tuple, Union, Mapping, Any, Callable, Iterable
+from typing import List, Tuple, Union, Mapping, Any, Callable, Iterable, TYPE_CHECKING
+
+# Plotting
+if find_spec("plotly") is None:
+    from bertopic._utils import MockPlotlyModule
+    plotting = MockPlotlyModule()
+
+else:
+    from bertopic import plotting
+    if TYPE_CHECKING:
+        import plotly.graph_objs as go
+        import matplotlib.figure as fig
+
 
 # Models
 try:
@@ -72,23 +86,9 @@
 )
 import bertopic._save_utils as save_utils
 
-
 logger = MyLogger()
 logger.configure("WARNING")
 
-try:
-    from bertopic import plotting
-    import plotly.graph_objects as go
-
-except ModuleNotFoundError as e:
-    if "No module named 'plotly'" in str(e):
-        logger.warning("Plotly is not installed. Please install it to use the plotting functions.")
-        from bertopic._utils import mock_plotly_go as go, MockPlotting
-
-        plotting = MockPlotting(logger)
-    else:
-        raise ModuleNotFoundError(e)
-
 
 class BERTopic:
     """BERTopic is a topic modeling technique that leverages BERT embeddings and
@@ -2415,7 +2415,7 @@ def visualize_topics(
         title: str = "<b>Intertopic Distance Map</b>",
         width: int = 650,
         height: int = 650,
-    ) -> go.Figure:
+    ) -> "go.Figure":
         """Visualize topics, their sizes, and their corresponding words.
 
         This visualization is highly inspired by LDAvis, a great visualization
@@ -2473,7 +2473,7 @@ def visualize_documents(
         title: str = "<b>Documents and Topics</b>",
         width: int = 1200,
         height: int = 750,
-    ) -> go.Figure:
+    ) -> "go.Figure":
         """Visualize documents and their topics in 2D.
 
         Arguments:
@@ -2575,7 +2575,7 @@ def visualize_document_datamap(
         topic_prefix: bool = False,
         datamap_kwds: dict = {},
         int_datamap_kwds: dict = {},
-    ):
+    ) -> "fig.Figure":
         """Visualize documents and their topics in 2D as a static plot for publication using
         DataMapPlot. This works best if there are between 5 and 60 topics. It is therefore best
         to use a sufficiently large `min_topic_size` or set `nr_topics` when building the model.
@@ -2686,7 +2686,7 @@ def visualize_hierarchical_documents(
         title: str = "<b>Hierarchical Documents and Topics</b>",
         width: int = 1200,
         height: int = 750,
-    ) -> go.Figure:
+    ) -> "go.Figure":
         """Visualize documents and their topics in 2D at different levels of hierarchy.
 
         Arguments:
@@ -2798,7 +2798,7 @@ def visualize_term_rank(
         title: str = "<b>Term score decline per Topic</b>",
         width: int = 800,
         height: int = 500,
-    ) -> go.Figure:
+    ) -> "go.Figure":
         """Visualize the ranks of all terms across all topics.
 
         Each topic is represented by a set of words. These words, however,
@@ -2863,7 +2863,7 @@ def visualize_topics_over_time(
         title: str = "<b>Topics over Time</b>",
         width: int = 1250,
         height: int = 450,
-    ) -> go.Figure:
+    ) -> "go.Figure":
         """Visualize topics over time.
 
         Arguments:
@@ -2919,7 +2919,7 @@ def visualize_topics_per_class(
         title: str = "<b>Topics per Class</b>",
         width: int = 1250,
         height: int = 900,
-    ) -> go.Figure:
+    ) -> "go.Figure":
         """Visualize topics per class.
 
         Arguments:
@@ -2973,7 +2973,7 @@ def visualize_distribution(
         title: str = "<b>Topic Probability Distribution</b>",
         width: int = 800,
         height: int = 600,
-    ) -> go.Figure:
+    ) -> "go.Figure":
         """Visualize the distribution of topic probabilities.
 
         Arguments:
@@ -3080,7 +3080,7 @@ def visualize_hierarchy(
         linkage_function: Callable[[csr_matrix], np.ndarray] = None,
         distance_function: Callable[[csr_matrix], csr_matrix] = None,
         color_threshold: int = 1,
-    ) -> go.Figure:
+    ) -> "go.Figure":
         """Visualize a hierarchical structure of the topics.
 
         A ward linkage function is used to perform the
@@ -3176,7 +3176,7 @@ def visualize_heatmap(
         title: str = "<b>Similarity Matrix</b>",
         width: int = 800,
         height: int = 800,
-    ) -> go.Figure:
+    ) -> "go.Figure":
         """Visualize a heatmap of the topic's similarity matrix.
 
         Based on the cosine similarity matrix between c-TF-IDFs or semantic embeddings of the topics,
@@ -3236,7 +3236,7 @@ def visualize_barchart(
         width: int = 250,
         height: int = 250,
         autoscale: bool = False,
-    ) -> go.Figure:
+    ) -> "go.Figure":
         """Visualize a barchart of selected topics.
 
         Arguments:

bertopic/_utils.py

@@ -4,7 +4,7 @@
 from collections.abc import Iterable
 from scipy.sparse import csr_matrix
 from scipy.spatial.distance import squareform
-from typing import Optional, Union, Tuple
+from typing import Optional, Union, Tuple, Any
 
 
 class MyLogger:
@@ -228,26 +228,11 @@ def to_ndarray(array: Union[np.ndarray, csr_matrix]) -> np.ndarray:
     return to_ndarray(repr_) if output_ndarray else repr_, ctfidf_used
 
 
-# Visualization mocks in case plotly is not installed
-class MockPlotting:
-    """Mock plotting module when plotly is not installed."""
+class MockPlotlyModule:
+    """Mock module that raises an error when plotly functions are called."""
 
-    def __init__(self, logger: MyLogger):
-        self.logger = logger
-
-    def __getattr__(self, name):
+    def __getattr__(self, name: str) -> Any:
         def mock_function(*args, **kwargs):
-            self.logger.warning(f"Plotly is not installed. Cannot use {name} visualization function.")
-            return MockFigure()
+            raise ImportError(f"Plotly is required to use '{name}'. " "Install it with uv pip install plotly")
 
         return mock_function
-
-
-class MockFigure:
-    """Mock class for plotly.graph_objects.Figure when plotly is not installed."""
-
-    def __init__(self, *args, **kwargs):
-        pass
-
-
-mock_plotly_go = type("MockPlotly", (), {"Figure": MockFigure})()

tests/test_other.py

@@ -2,13 +2,9 @@
 from bertopic.dimensionality import BaseDimensionalityReduction
 
 try:
-    import plotly.graph_objects as go
-
-    figure_type = go.Figure
+    from plotly.graph_objects import Figure
 except ImportError:
-    from bertopic._utils import MockFigure
-
-    figure_type = MockFigure
+    Figure = None
 
 
 def test_load_save_model():
@@ -41,5 +37,9 @@ def test_no_plotly():
         umap_model=BaseDimensionalityReduction(),
     )
     model.fit(["hello", "hi", "goodbye", "goodbye", "whats up"] * 10)
-    out = model.visualize_topics()
-    assert isinstance(out, figure_type)
+    
+    try:
+        out = model.visualize_topics()
+        assert isinstance(out, Figure) if Figure  else False
+    except ImportError as e:
+        assert "Plotly is required to use" in str(e)