feat: support hightlighing except for segmented plots and boxplots (#59)

Author: SaaiVenkat

example/line/matplotlib/example_mpl_line.py

@@ -27,4 +27,3 @@
 # Show the plot
 plt.show()
 maidr.show(line_plot)
-print()

maidr/__init__.py

@@ -2,7 +2,16 @@
 
 from .core import Maidr
 from .core.enum import PlotType
-from .patch import barplot, boxplot, clear, heatmap, histogram, lineplot, scatterplot
+from .patch import (
+    barplot,
+    boxplot,
+    clear,
+    heatmap,
+    highlight,
+    histogram,
+    lineplot,
+    scatterplot,
+)
 from .maidr import close, save_html, show, stacked
 
 __all__ = [

maidr/core/context_manager.py

@@ -1,9 +1,13 @@
+from __future__ import annotations
+
 import contextlib
 import contextvars
 import threading
 
 import wrapt
 
+from maidr.core.plot.boxplot import BoxPlotContainer
+
 
 class ContextManager:
     _instance = None
@@ -18,6 +22,10 @@ def __new__(cls):
                     cls._instance = super(ContextManager, cls).__new__()
         return cls._instance
 
+    @classmethod
+    def is_internal_context(cls):
+        return cls._internal_context.get()
+
     @classmethod
     @contextlib.contextmanager
     def set_internal_context(cls):
@@ -27,10 +35,6 @@ def set_internal_context(cls):
         finally:
             cls._internal_context.reset(token_internal_context)
 
-    @classmethod
-    def is_internal_context(cls):
-        return cls._internal_context.get()
-
 
 @wrapt.decorator
 def manage_context(wrapped=None, _=None, args=None, kwargs=None):
@@ -41,3 +45,74 @@ def manage_context(wrapped=None, _=None, args=None, kwargs=None):
     # Set the internal context to avoid cyclic processing.
     with ContextManager.set_internal_context():
         return wrapped(*args, **kwargs)
+
+
+class BoxplotContextManager(ContextManager):
+    _bxp_context = contextvars.ContextVar("bxp_context", default=BoxPlotContainer())
+
+    @classmethod
+    @contextlib.contextmanager
+    def set_internal_context(cls):
+        with super(BoxplotContextManager, cls).set_internal_context():
+            token = cls._bxp_context.set(BoxPlotContainer())
+            try:
+                yield cls.get_bxp_context()
+            finally:
+                cls._bxp_context.reset(token)
+
+    @classmethod
+    def get_bxp_context(cls) -> BoxPlotContainer:
+        return cls._bxp_context.get()
+
+    @classmethod
+    def add_bxp_context(cls, bxp_context: dict) -> None:
+        cls.get_bxp_context().add_artists(bxp_context)
+
+    @classmethod
+    def set_bxp_orientation(cls, orientation: str) -> None:
+        cls.get_bxp_context().set_orientation(orientation)
+
+
+class HighlightContextManager:
+    _instance = None
+    _lock = threading.Lock()
+
+    _maidr_element = contextvars.ContextVar("_maidr_element", default=False)
+    _elements = contextvars.ContextVar("elements", default=[])
+
+    def __new__(cls):
+        if not cls._instance:
+            with cls._lock:
+                if not cls._instance:
+                    cls._instance = super(HighlightContextManager, cls).__new__()
+        return cls._instance
+
+    @classmethod
+    def is_maidr_element(cls):
+        return cls._maidr_element.get()
+
+    @classmethod
+    @contextlib.contextmanager
+    def set_maidr_element(cls, element):
+        if element not in cls._elements.get():
+            yield
+            return
+
+        token_maidr_element = cls._maidr_element.set(True)
+        try:
+            yield
+        finally:
+            cls._maidr_element.reset(token_maidr_element)
+            # Remove element from the context list after tagging
+            new_elements = cls._elements.get().copy()
+            new_elements.remove(element)
+            cls._elements.set(new_elements)
+
+    @classmethod
+    @contextlib.contextmanager
+    def set_maidr_elements(cls, elements: list):
+        token_paths = cls._elements.set(elements)
+        try:
+            yield
+        finally:
+            cls._elements.reset(token_paths)

maidr/core/maidr.py

@@ -10,6 +10,7 @@
 
 from matplotlib.figure import Figure
 
+from maidr.core.context_manager import HighlightContextManager
 from maidr.core.plot import MaidrPlot
 
 
@@ -104,7 +105,9 @@ def destroy(self) -> None:
 
     def _create_html_tag(self) -> Tag:
         """Create the MAIDR HTML using HTML tags."""
-        svg = self._get_svg()
+        tagged_elements = [element for plot in self._plots for element in plot.elements]
+        with HighlightContextManager.set_maidr_elements(tagged_elements):
+            svg = self._get_svg()
         maidr = f"\nlet maidr = {json.dumps(self._flatten_maidr(), indent=2)}\n"
 
         # Inject plot's svg and MAIDR structure into html tag.

maidr/core/plot/barplot.py

@@ -51,4 +51,7 @@ def _extract_bar_container_data(
         if len(plot) != len(level):
             return None
 
+        # Tag the elements for highlighting.
+        self._elements.extend(plot)
+
         return [float(patch.get_height()) for patch in plot]

maidr/core/plot/boxplot.py

@@ -126,10 +126,12 @@ class BoxPlot(
     DictMergerMixin,
 ):
     def __init__(self, ax: Axes, **kwargs) -> None:
+        super().__init__(ax, PlotType.BOX)
+
         self._bxp_stats = kwargs.pop("bxp_stats", None)
         self._orientation = kwargs.pop("orientation", "vert")
         self._bxp_extractor = BoxPlotExtractor(orientation=self._orientation)
-        super().__init__(ax, PlotType.BOX)
+        self._support_highlighting = False
 
     def render(self) -> dict:
         base_schema = super().render()

maidr/core/plot/grouped_barplot.py

@@ -18,6 +18,7 @@ class GroupedBarPlot(
 ):
     def __init__(self, ax: Axes, plot_type: PlotType) -> None:
         super().__init__(ax, plot_type)
+        self._support_highlighting = False
 
     def _extract_axes_data(self) -> dict:
         base_ax_schema = super()._extract_axes_data()

maidr/core/plot/heatmap.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import numpy.ma as ma
+
 from matplotlib.axes import Axes
 from matplotlib.cm import ScalarMappable
 from matplotlib.collections import QuadMesh
@@ -44,23 +46,29 @@ def _extract_axes_data(self) -> dict:
 
     def _extract_plot_data(self) -> list[list]:
         plot = self.extract_scalar_mappable(self.ax)
-        data = HeatPlot._extract_scalar_mappable_data(plot)
+        data = self._extract_scalar_mappable_data(plot)
 
         if data is None:
             raise ExtractionError(self.type, plot)
 
         return data
 
-    @staticmethod
-    def _extract_scalar_mappable_data(sm: ScalarMappable | None) -> list[list] | None:
+    def _extract_scalar_mappable_data(
+        self, sm: ScalarMappable | None
+    ) -> list[list] | None:
         if sm is None or sm.get_array() is None:
             return None
 
         array = sm.get_array().data
         if isinstance(sm, QuadMesh):
             # Data gets flattened in ScalarMappable, when the plot is from QuadMesh.
             # So, reshaping the data to reflect the original quadrilaterals
-            m, n, _ = sm.get_coordinates().shape
+            m, n, _ = ma.shape(sm.get_coordinates())
             array = array.reshape(m - 1, n - 1)  # Coordinates shape is (M + 1, N + 1)
 
+            # Tag the elements for highlighting
+            self._elements.append(sm)
+        else:
+            self._support_highlighting = False
+
         return [list(map(float, row)) for row in array]

maidr/core/plot/histogram.py

@@ -15,15 +15,16 @@ def __init__(self, ax: Axes) -> None:
 
     def _extract_plot_data(self) -> list[dict]:
         plot = self.extract_container(self.ax, BarContainer)
-        data = HistPlot._extract_bar_container_data(plot)
+        data = self._extract_bar_container_data(plot)
 
         if data is None:
             raise ExtractionError(self.type, plot)
 
         return data
 
-    @staticmethod
-    def _extract_bar_container_data(plot: BarContainer | None) -> list[dict] | None:
+    def _extract_bar_container_data(
+        self, plot: BarContainer | None
+    ) -> list[dict] | None:
         if plot is None or plot.patches is None:
             return None
 
@@ -44,4 +45,7 @@ def _extract_bar_container_data(plot: BarContainer | None) -> list[dict] | None:
                 }
             )
 
+        # Tag the elements for highlighting
+        self._elements.extend(plot.patches)
+
         return data

maidr/core/plot/lineplot.py

@@ -13,24 +13,29 @@ class LinePlot(MaidrPlot, LineExtractorMixin):
     def __init__(self, ax: Axes) -> None:
         super().__init__(ax, PlotType.LINE)
 
+    def _get_selector(self) -> str:
+        return "g[maidr='true'] > path"
+
     def _extract_plot_data(self) -> list[dict]:
         plot = self.extract_line(self.ax)
-        data = LinePlot._extract_line_data(plot)
+        data = self._extract_line_data(plot)
 
         if data is None:
             raise ExtractionError(self.type, plot)
 
         return data
 
-    @staticmethod
-    def _extract_line_data(plot: Line2D | None) -> list[dict] | None:
+    def _extract_line_data(self, plot: Line2D | None) -> list[dict] | None:
         if plot is None or plot.get_xydata() is None:
             return None
 
+        # Tag the elements for highlighting.
+        self._elements.append(plot)
+
         return [
             {
-                MaidrKey.X.value: float(x),
-                MaidrKey.Y.value: float(y),
+                MaidrKey.X: float(x),
+                MaidrKey.Y: float(y),
             }
             for x, y in plot.get_xydata()
         ]

maidr/core/plot/maidr_plot.py

@@ -37,20 +37,32 @@ class MaidrPlot(ABC):
     def __init__(self, ax: Axes, plot_type: PlotType) -> None:
         # graphic object
         self.ax = ax
+        self._support_highlighting = True
+        self._elements = []
 
         # MAIDR data
         self.type = plot_type
         self._schema = {}
 
     def render(self) -> dict:
         """Initialize the MAIDR schema dictionary with basic plot information."""
-        return {
+        maidr_schema = {
             MaidrKey.TYPE: self.type,
             MaidrKey.TITLE: self.ax.get_title(),
             MaidrKey.AXES: self._extract_axes_data(),
             MaidrKey.DATA: self._extract_plot_data(),
         }
 
+        # Include selector only if the plot supports highlighting.
+        if self._support_highlighting:
+            maidr_schema[MaidrKey.SELECTOR] = self._get_selector()
+
+        return maidr_schema
+
+    def _get_selector(self) -> str:
+        """Return the CSS selector for highlighting elements."""
+        return "path[maidr='true']"
+
     def _extract_axes_data(self) -> dict:
         """Extract the plot's axes data"""
         return {
@@ -74,8 +86,14 @@ def schema(self) -> dict:
             self._schema = self.render()
         return self._schema
 
+    @property
+    def elements(self) -> list:
+        if not self._schema:
+            self._schema = self.render()
+        return self._elements
+
     def set_id(self, maidr_id: str) -> None:
         """Set the unique identifier for the plot within the MAIDR schema."""
         if not self._schema:
             self._schema = self.render()
-        self._schema[MaidrKey.ID.value] = maidr_id
+        self._schema[MaidrKey.ID] = maidr_id

maidr/core/plot/scatterplot.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import numpy.ma as ma
+
 from matplotlib.axes import Axes
 from matplotlib.collections import PathCollection
 
@@ -13,28 +15,29 @@ class ScatterPlot(MaidrPlot, CollectionExtractorMixin):
     def __init__(self, ax: Axes) -> None:
         super().__init__(ax, PlotType.SCATTER)
 
+    def _get_selector(self) -> str:
+        return "g[maidr='true'] > use"
+
     def _extract_plot_data(self) -> list[dict]:
         plot = self.extract_collection(self.ax, PathCollection)
-        data = ScatterPlot._extract_point_data(plot)
+        data = self._extract_point_data(plot)
 
         if data is None:
             raise ExtractionError(self.type, plot)
 
         return data
 
-    @staticmethod
-    def _extract_point_data(plot: PathCollection | None) -> list[dict] | None:
+    def _extract_point_data(self, plot: PathCollection | None) -> list[dict] | None:
         if plot is None or plot.get_offsets() is None:
             return None
 
-        data = []
-        for point in plot.get_offsets().data:
-            x, y = point
-            data.append(
-                {
-                    MaidrKey.X.value: float(x),
-                    MaidrKey.Y.value: float(y),
-                }
-            )
+        # Tag the elements for highlighting.
+        self._elements.append(plot)
 
-        return data
+        return [
+            {
+                MaidrKey.X: float(x),
+                MaidrKey.Y: float(y),
+            }
+            for x, y in ma.getdata(plot.get_offsets())
+        ]