maidr 1.3.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

maidr/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.0"
+__version__ = "1.4.0"
 
 from .api import close, render, save_html, show, stacked
 from .core import Maidr
@@ -15,6 +15,7 @@ from .patch import (
     scatterplot,
     regplot,
     kdeplot,
+    mplfinance,
 )
 
 __all__ = [

maidr/core/plot/candlestick.py

@@ -1,11 +1,12 @@
 import matplotlib.dates as mdates
 import numpy as np
 from matplotlib.axes import Axes
-from matplotlib.collections import LineCollection, PatchCollection
 from matplotlib.patches import Rectangle
 
 from maidr.core.enum.plot_type import PlotType
 from maidr.core.plot.maidr_plot import MaidrPlot
+from maidr.core.enum.maidr_key import MaidrKey
+from maidr.util.mplfinance_utils import MplfinanceDataExtractor
 
 
 class CandlestickPlot(MaidrPlot):
@@ -27,213 +28,96 @@ class CandlestickPlot(MaidrPlot):
             raise ValueError("Axes list cannot be empty.")
         super().__init__(axes[0], PlotType.CANDLESTICK)
 
+        # Store custom collections passed from mplfinance patch
+        self._maidr_wick_collection = kwargs.get("_maidr_wick_collection", None)
+        self._maidr_body_collection = kwargs.get("_maidr_body_collection", None)
+        self._maidr_date_nums = kwargs.get("_maidr_date_nums", None)
+
+        # Store the GID for proper selector generation
+        self._maidr_gid = None
+        if self._maidr_body_collection:
+            self._maidr_gid = self._maidr_body_collection.get_gid()
+        elif self._maidr_wick_collection:
+            self._maidr_gid = self._maidr_wick_collection.get_gid()
+
     def _extract_plot_data(self) -> list[dict]:
-        """
-        Extracts candlestick (OHLC) and volume data from the plot axes.
-
-        This method assumes that the candlestick chart is structured with
-        LineCollection for wicks and PatchCollection of Rectangles for bodies
-        on the first axis (self.axes[0]). Volume data is expected as a
-        PatchCollection of Rectangles on the second axis (self.axes[1]), if present.
-        Open and close prices are inferred from the body rectangle's color.
-
-        Returns
-        -------
-        list[dict]
-            A list of dictionaries, where each dictionary represents a data point
-            with 'value' (date string YYYY-MM-DD), 'open', 'high', 'low',
-            'close', and 'volume'. Fields that cannot be extracted will be None.
-
-        Examples
-        --------
-        Assuming a plot has been generated and `plot_instance.axes` is populated:
-        >>> data = plot_instance._extract_plot_data()
-        >>> print(data[0])
-        {
-            'value': '2021-01-01',
-            'open': 100.0,
-            'high': 100.9,
-            'low': 99.27,
-            'close': 100.75,
-            'volume': 171914,
-        }
-        """
+        """Extract candlestick data from the plot."""
+
+        # Get the custom collections from kwargs
+        body_collection = self._maidr_body_collection
+        wick_collection = self._maidr_wick_collection
+
+        if body_collection and wick_collection:
+            # Store the GID from the body collection for highlighting
+            self._maidr_gid = body_collection.get_gid()
+
+            # Use the original collections for highlighting
+            self._elements = [body_collection, wick_collection]
+
+            # Use the utility class to extract data
+            return MplfinanceDataExtractor.extract_candlestick_data(
+                body_collection, wick_collection, self._maidr_date_nums
+            )
+
+        # Fallback to original detection method
         if not self.axes:
             return []
 
-        plot_data: list[dict] = []
-        ax_ohlc: Axes = self.axes[0]
-
-        body_rectangles: list[Rectangle] = []
-        wick_collection: LineCollection | None = None
-
-        # Find candlestick body Rectangles from the OHLC axis
-        # Prefer PatchCollection containing Rectangles, fallback to individual Rectangles in ax.patches
-        for collection in ax_ohlc.collections:
-            if isinstance(collection, PatchCollection):
-                # Check if the collection's patches are Rectangles
-                try:
-                    # Iterating a PatchCollection yields its constituent Patch objects
-                    patches_are_rects = all(
-                        isinstance(p, Rectangle) for p in collection
-                    )
-                    if (
-                        patches_are_rects and len(collection.get_paths()) > 0
-                    ):  # Ensure it has paths and they are Rectangles
-                        for (
-                            patch
-                        ) in collection:  # Iterate to get actual Rectangle objects
-                            if isinstance(patch, Rectangle):
-                                body_rectangles.append(patch)
-                        if (
-                            body_rectangles
-                        ):  # If we found rectangles this way, assume this is the primary body collection
-                            break
-                except Exception:
-                    # Could fail if collection is not iterable in the expected way or patches are not Rectangles
-                    pass
-
-        if not body_rectangles:
-            for patch in ax_ohlc.patches:
-                if isinstance(patch, Rectangle):
-                    body_rectangles.append(patch)
-
-        if not body_rectangles:
-            pass
-
-        ax_for_wicks: Axes | None = None
-        if len(self.axes) > 1:
-            ax_for_wicks = self.axes[1]
-
-        if ax_for_wicks:
-            # Attempt 1: Find wicks in ax_for_wicks.collections (as a LineCollection)
-            for collection in ax_for_wicks.collections:
-                if isinstance(collection, LineCollection):
-                    segments = collection.get_segments()
-                    # Check if the collection contains segments and the first segment looks like a vertical line
-                    if segments is not None and len(segments) > 0:
-                        first_segment = segments[0]
-                        if (
-                            len(first_segment) == 2  # Segment consists of two points
-                            and len(first_segment[0]) == 2  # First point has (x, y)
-                            and len(first_segment[1]) == 2  # Second point has (x, y)
-                            and np.isclose(
-                                first_segment[0][0], first_segment[1][0]
-                            )  # X-coordinates are close (vertical)
-                        ):
-                            wick_collection = collection
-                            break  # Found a suitable LineCollection
-
-            # Attempt 2: If no LineCollection found, try to find wicks from individual Line2D objects in ax_for_wicks.get_lines()
-            if not wick_collection and hasattr(ax_for_wicks, "get_lines"):
-                potential_wick_segments = []
-                for line in ax_for_wicks.get_lines():  # Iterate over Line2D objects
-                    x_data, y_data = line.get_data()
-                    # A wick is typically a vertical line defined by two points.
-                    if len(x_data) == 2 and len(y_data) == 2:
-                        if np.isclose(x_data[0], x_data[1]):  # Check for verticality
-                            # Create a segment in the format [[x1, y1], [x2, y2]]
-                            segment = [
-                                [x_data[0], y_data[0]],
-                                [x_data[1], y_data[1]],
-                            ]
-                            potential_wick_segments.append(segment)
-
-                if potential_wick_segments:
-                    # If wick segments were found from individual lines,
-                    # create a new LineCollection to hold them.
-                    # This allows the downstream processing logic
-                    # for wicks to remain consistent.
-                    # Basic properties are set; color/linestyle
-                    # are defaults and may not match
-                    # the original plot's styling if that
-                    # were relevant for segment extraction.
-                    wick_collection = LineCollection(
-                        potential_wick_segments,
-                        colors="k",  # Default color for the temporary collection
-                        linestyles="solid",  # Default linestyle
-                    )
-
-        # Process wicks into a map: x_coordinate -> (low_price, high_price)
-        wick_segments_map: dict[float, tuple[float, float]] = {}
-        if wick_collection:
-            for seg in wick_collection.get_segments():
-                if len(seg) == 2 and len(seg[0]) == 2 and len(seg[1]) == 2:
-                    # Ensure x-coordinates are (nearly) identical for a vertical wick line
-                    if np.isclose(seg[0][0], seg[1][0]):
-                        x_coord = seg[0][0]  # Matplotlib date number
-                        low_price = min(seg[0][1], seg[1][1])
-                        high_price = max(seg[0][1], seg[1][1])
-                        wick_segments_map[x_coord] = (low_price, high_price)
-
-        body_rectangles.sort(key=lambda r: r.get_x())
-
-        for rect in body_rectangles:
-            x_left = rect.get_x()
-            width = rect.get_width()
-            x_center_num = x_left + width / 2.0
-
-            try:
-                date_dt = mdates.num2date(x_center_num)
-                date_str = date_dt.strftime("%Y-%m-%d")
-            except ValueError:
-                date_str = f"raw_date_{x_center_num:.2f}"
-
-            y_bottom = rect.get_y()
-            height = rect.get_height()
-            face_color = rect.get_facecolor()  # RGBA tuple
-
-            # Infer open and close prices
-            # Heuristic: Green component > Red component for an "up" candle (close > open)
-            # This assumes standard green for up, red for down.
-            # A more robust method would involve knowing the exact up/down colors used.
-            is_up_candle = (
-                face_color[1] > face_color[0]
-            )  # Compare Green and Red components
-
-            if is_up_candle:  # Typically green: price went up
-                open_price = y_bottom
-                close_price = y_bottom + height
-            else:  # Typically red: price went down (or other color)
-                close_price = y_bottom
-                open_price = y_bottom + height
-
-            matched_wick_data = None
-            closest_wick_x = None
-            min_diff = float("inf")
-
-            for wick_x, prices in wick_segments_map.items():
-                diff = abs(wick_x - x_center_num)
-                if diff < min_diff:
-                    min_diff = diff
-                    closest_wick_x = wick_x
-
-            # Tolerance for matching wick x-coordinate (e.g., 10% of candle width)
-            if closest_wick_x is not None and min_diff < (width * 0.1):
-                matched_wick_data = wick_segments_map[closest_wick_x]
-
-            if matched_wick_data:
-                low_price, high_price = matched_wick_data
-                # Ensure high >= max(open,close) and low <= min(open,close)
-                high_price = max(high_price, open_price, close_price)
-                low_price = min(low_price, open_price, close_price)
-            else:
-                # Fallback if no wick found: high is max(open,close), low is min(open,close)
-                high_price = max(open_price, close_price)
-                low_price = min(open_price, close_price)
-
-            plot_data.append(
-                {
-                    "value": date_str,
-                    "open": open_price,
-                    "high": high_price,
-                    "low": low_price,
-                    "close": close_price,
-                    "volume": 0,
-                }
+        ax_ohlc = self.axes[0]
+        candles = []
+
+        # Look for Rectangle patches (original_flavor candlestick)
+        body_rectangles = []
+        for patch in ax_ohlc.patches:
+            if isinstance(patch, Rectangle):
+                body_rectangles.append(patch)
+
+        if body_rectangles:
+            # Set elements for highlighting
+            self._elements = body_rectangles
+
+            # Generate a GID for highlighting if none exists
+            if not self._maidr_gid:
+                import uuid
+
+                self._maidr_gid = f"maidr-{uuid.uuid4()}"
+                # Set GID on all rectangles
+                for rect in body_rectangles:
+                    rect.set_gid(self._maidr_gid)
+
+            # Use the utility class to extract data
+            return MplfinanceDataExtractor.extract_rectangle_candlestick_data(
+                body_rectangles, self._maidr_date_nums
             )
-        self._elements.extend(body_rectangles)
-        return plot_data
+
+        return []
 
     def _extract_axes_data(self) -> dict:
         return {}
+
+    def _get_selector(self) -> str:
+        """Return the CSS selector for highlighting candlestick elements in the SVG output."""
+        # Use the stored GID if available, otherwise fall back to generic selector
+        if self._maidr_gid:
+            # Use the full GID as the id attribute (since that's what's in the SVG)
+            selector = (
+                f"g[id='{self._maidr_gid}'] > path, g[id='{self._maidr_gid}'] > rect"
+            )
+        else:
+            selector = "g[maidr='true'] > path, g[maidr='true'] > rect"
+        return selector
+
+    def render(self) -> dict:
+        """Initialize the MAIDR schema dictionary with basic plot information."""
+        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

maidr/core/plot/maidr_plot.py

@@ -42,8 +42,13 @@ class MaidrPlot(ABC):
         self._support_highlighting = True
         self._elements = []
         ss = self.ax.get_subplotspec()
-        self.row_index = ss.rowspan.start
-        self.col_index = ss.colspan.start
+        # Handle cases where subplotspec is None (dynamically created axes)
+        if ss is not None:
+            self.row_index = ss.rowspan.start
+            self.col_index = ss.colspan.start
+        else:
+            self.row_index = 0
+            self.col_index = 0
 
         # MAIDR data
         self.type = plot_type

maidr/core/plot/maidr_plot_factory.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from matplotlib.axes import Axes
-
 from maidr.core.enum import PlotType
 from maidr.core.plot.barplot import BarPlot
 from maidr.core.plot.boxplot import BoxPlot
@@ -13,6 +12,9 @@ from maidr.core.plot.lineplot import MultiLinePlot
 from maidr.core.plot.maidr_plot import MaidrPlot
 from maidr.core.plot.scatterplot import ScatterPlot
 from maidr.core.plot.regplot import SmoothPlot
+from maidr.core.plot.mplfinance_barplot import MplfinanceBarPlot
+from maidr.core.plot.mplfinance_lineplot import MplfinanceLinePlot
+from maidr.util.plot_detection import PlotDetectionUtils
 
 
 class MaidrPlotFactory:
@@ -38,17 +40,13 @@ class MaidrPlotFactory:
             single_ax = ax
 
         if plot_type == PlotType.CANDLESTICK:
-            if isinstance(ax, list):
-                # If ax is a list of lists, flatten it
-                if ax and isinstance(ax[0], list):
-                    axes = ax[0]  # Take the first inner list
-                else:
-                    axes = ax  # Use the list as-is
-            else:
-                axes = [ax]  # Wrap single axes in list
+            axes = PlotDetectionUtils.get_candlestick_axes(ax)
             return CandlestickPlot(axes, **kwargs)
         elif PlotType.BAR == plot_type or PlotType.COUNT == plot_type:
-            return BarPlot(single_ax)
+            if PlotDetectionUtils.is_mplfinance_bar_plot(**kwargs):
+                return MplfinanceBarPlot(single_ax, **kwargs)
+            else:
+                return BarPlot(single_ax)
         elif PlotType.BOX == plot_type:
             return BoxPlot(single_ax, **kwargs)
         elif PlotType.HEAT == plot_type:
@@ -56,7 +54,10 @@ class MaidrPlotFactory:
         elif PlotType.HIST == plot_type:
             return HistPlot(single_ax)
         elif PlotType.LINE == plot_type:
-            return MultiLinePlot(single_ax)
+            if PlotDetectionUtils.is_mplfinance_line_plot(single_ax, **kwargs):
+                return MplfinanceLinePlot(single_ax, **kwargs)
+            else:
+                return MultiLinePlot(single_ax, **kwargs)
         elif PlotType.SCATTER == plot_type:
             return ScatterPlot(single_ax)
         elif PlotType.DODGED == plot_type or PlotType.STACKED == plot_type:

maidr/core/plot/mplfinance_barplot.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+from matplotlib.axes import Axes
+from matplotlib.patches import Rectangle
+
+from maidr.core.enum import PlotType
+from maidr.core.plot import MaidrPlot
+from maidr.exception import ExtractionError
+from maidr.util.mixin import (
+    ContainerExtractorMixin,
+    DictMergerMixin,
+    LevelExtractorMixin,
+)
+from maidr.util.mplfinance_utils import MplfinanceDataExtractor
+
+
+class MplfinanceBarPlot(
+    MaidrPlot, ContainerExtractorMixin, LevelExtractorMixin, DictMergerMixin
+):
+    """
+    Specialized bar plot class for mplfinance volume bars.
+
+    This class handles the extraction and processing of volume data from mplfinance
+    plots, including proper date conversion and data validation.
+    """
+
+    def __init__(self, ax: Axes, **kwargs) -> None:
+        super().__init__(ax, PlotType.BAR)
+        # Store custom patches passed from mplfinance patch
+        self._custom_patches = kwargs.get("_maidr_patches", None)
+        # Store date numbers for volume bars (from mplfinance)
+        self._maidr_date_nums = kwargs.get("_maidr_date_nums", None)
+
+    def _extract_plot_data(self) -> list:
+        """Extract data from mplfinance volume patches."""
+        if self._custom_patches:
+            return self._extract_volume_patches_data(self._custom_patches)
+
+        # Fallback to original bar plot logic if no custom patches
+        plot = self.extract_container(
+            self.ax, ContainerExtractorMixin, include_all=True
+        )
+        data = self._extract_bar_container_data(plot)
+        levels = self.extract_level(self.ax)
+        formatted_data = []
+        combined_data = list(
+            zip(levels, data) if plot[0].orientation == "vertical" else zip(data, levels)  # type: ignore
+        )
+        if combined_data:  # type: ignore
+            for x, y in combined_data:  # type: ignore
+                formatted_data.append({"x": x, "y": y})
+            return formatted_data
+        if len(formatted_data) == 0:
+            raise ExtractionError(self.type, plot)
+        if data is None:
+            raise ExtractionError(self.type, plot)
+
+        return data
+
+    def _extract_volume_patches_data(self, volume_patches: list[Rectangle]) -> list:
+        """
+        Extract data from volume Rectangle patches (used by mplfinance volume bars).
+        """
+        if not volume_patches:
+            return []
+
+        # Sort patches by x-coordinate to maintain order
+        sorted_patches = sorted(volume_patches, key=lambda p: p.get_x())
+
+        # Set elements for highlighting (use the patches directly)
+        self._elements = sorted_patches
+
+        # Use the utility class to extract data
+        return MplfinanceDataExtractor.extract_volume_data(
+            sorted_patches, self._maidr_date_nums
+        )
+
+    def _extract_bar_container_data(self, plot: list | None) -> list | None:
+        """Fallback method for regular bar containers."""
+        if plot is None:
+            return None
+
+        # Since v0.13, Seaborn has transitioned from using `list[Patch]` to
+        # `list[BarContainers] for plotting bar plots.
+        # So, extract data correspondingly based on the level.
+        # Flatten all the `list[BarContainer]` to `list[Patch]`.
+        patches = [patch for container in plot for patch in container.patches]
+        level = self.extract_level(self.ax)
+        if level is None or len(level) == 0:  # type: ignore
+            level = ["" for _ in range(len(patches))]  # type: ignore
+
+        if len(patches) != len(level):
+            return None
+
+        self._elements.extend(patches)
+
+        return [float(patch.get_height()) for patch in patches]
+
+    def _extract_axes_data(self) -> dict:
+        """Extract axes data with mplfinance-specific cleaning."""
+        ax_data = super()._extract_axes_data()
+
+        # For mplfinance volume plots, clean up the y-axis label
+        if self._custom_patches:
+            y_label = ax_data.get("y")
+            if y_label:
+                ax_data["y"] = MplfinanceDataExtractor.clean_axis_label(y_label)
+
+        return ax_data
+
+    def _get_selector(self) -> str:
+        """Return the CSS selector for highlighting bar elements in the SVG output."""
+        # Use the standard working selector that gets replaced with UUID by Maidr class
+        # This works for both original bar plots and mplfinance volume bars
+        return "g[maidr='true'] > path"

maidr/core/plot/mplfinance_lineplot.py

@@ -0,0 +1,146 @@
+from typing import List, Union
+
+from matplotlib.axes import Axes
+from matplotlib.lines import Line2D
+import numpy as np
+
+from maidr.core.enum.maidr_key import MaidrKey
+from maidr.core.enum.plot_type import PlotType
+from maidr.core.plot.maidr_plot import MaidrPlot
+from maidr.exception.extraction_error import ExtractionError
+from maidr.util.mixin.extractor_mixin import LineExtractorMixin
+from maidr.util.mplfinance_utils import MplfinanceDataExtractor
+import uuid
+
+
+class MplfinanceLinePlot(MaidrPlot, LineExtractorMixin):
+    """
+    Specialized line plot class for mplfinance moving averages.
+
+    This class handles the extraction and processing of moving average data from mplfinance
+    plots, including proper date conversion, NaN filtering, and moving average period detection.
+    """
+
+    def __init__(self, ax: Axes, **kwargs):
+        super().__init__(ax, PlotType.LINE)
+
+    def _get_selector(self) -> Union[str, List[str]]:
+        """Return selectors for all lines that have data."""
+        all_lines = self.ax.get_lines()
+        if not all_lines:
+            return ["g[maidr='true'] > path"]
+
+        selectors = []
+        for line in all_lines:
+            # Only create selectors for lines that have data
+            xydata = line.get_xydata()
+            if xydata is None or not xydata.size:  # type: ignore
+                continue
+            gid = line.get_gid()
+            if gid:
+                selectors.append(f"g[id='{gid}'] path")
+            else:
+                selectors.append("g[maidr='true'] > path")
+
+        if not selectors:
+            return ["g[maidr='true'] > path"]
+
+        return selectors
+
+    def _extract_plot_data(self) -> Union[List[List[dict]], None]:
+        """Extract data from mplfinance moving average lines."""
+        data = self._extract_line_data()
+
+        if data is None:
+            raise ExtractionError(self.type, None)
+
+        return data
+
+    def _extract_line_data(self) -> Union[List[List[dict]], None]:
+        """
+        Extract data from all line objects and return as separate arrays.
+
+        This method handles mplfinance-specific logic including:
+        - Date conversion from matplotlib date numbers
+        - NaN filtering for moving averages
+        - Moving average period detection
+        - Proper data validation
+
+        Returns
+        -------
+        list[list[dict]] | None
+            List of lists, where each inner list contains dictionaries with x,y coordinates
+            and line identifiers for one line, or None if the plot data is invalid.
+        """
+        all_lines = self.ax.get_lines()
+        if not all_lines:
+            return None
+
+        all_lines_data = []
+
+        for line_idx, line in enumerate(all_lines):
+            xydata = line.get_xydata()
+            if xydata is None or not xydata.size:  # type: ignore
+                continue
+
+            self._elements.append(line)
+
+            # Assign unique GID to each line if not already set
+            if line.get_gid() is None:
+                unique_gid = f"maidr-{uuid.uuid4()}"
+                line.set_gid(unique_gid)
+
+            label: str = line.get_label()  # type: ignore
+            line_data = []
+
+            # Check if this line has date numbers from mplfinance
+            date_nums = getattr(line, "_maidr_date_nums", None)
+
+            # Convert xydata to list of points
+            for i, (x, y) in enumerate(line.get_xydata()):  # type: ignore
+                # Skip points with NaN or inf values to prevent JSON parsing errors
+                if np.isnan(x) or np.isnan(y) or np.isinf(x) or np.isinf(y):
+                    continue
+
+                # Handle x-value conversion - could be string (date) or numeric
+                if isinstance(x, str):
+                    x_value = x  # Keep string as-is (for dates)
+                else:
+                    # Check if we have date numbers from mplfinance
+                    if date_nums is not None and i < len(date_nums):
+                        # Use the date number to convert to date string
+                        date_num = float(date_nums[i])
+                        x_value = self._convert_x_to_date(date_num)
+                    else:
+                        x_value = float(x)  # Convert numeric to float
+
+                point_data = {
+                    MaidrKey.X: x_value,
+                    MaidrKey.Y: float(y),
+                    MaidrKey.FILL: (label if not label.startswith("_child") else ""),
+                }
+                line_data.append(point_data)
+
+            if line_data:
+                all_lines_data.append(line_data)
+
+        return all_lines_data if all_lines_data else None
+
+    def _convert_x_to_date(self, x_value: float) -> str:
+        """
+        Convert x-coordinate to date string for mplfinance plots.
+
+        This method uses the MplfinanceDataExtractor utility to convert
+        matplotlib date numbers to proper date strings.
+
+        Parameters
+        ----------
+        x_value : float
+            The x-coordinate value (matplotlib date number)
+
+        Returns
+        -------
+        str
+            Date string in YYYY-MM-DD format
+        """
+        return MplfinanceDataExtractor._convert_date_num_to_string(x_value)

maidr/patch/__init__.py

@@ -0,0 +1,15 @@
+# Import all patches to ensure they are applied
+from . import (
+    barplot,
+    boxplot,
+    candlestick,
+    clear,
+    heatmap,
+    highlight,
+    histogram,
+    lineplot,
+    scatterplot,
+    regplot,
+    kdeplot,
+    mplfinance,
+)