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

maidr/core/plot/mplfinance_barplot.py

@@ -0,0 +1,139 @@
+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)
+        # Store custom title
+
+    def set_title(self, title: str) -> None:
+        """Set a custom title for this volume bar plot."""
+        self._title = title
+
+    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"
+
+    def render(self) -> dict:
+        """Initialize the MAIDR schema dictionary with basic plot information."""
+        from maidr.core.enum.maidr_key import MaidrKey
+
+        title = "Volume Bar Plot"
+
+        maidr_schema = {
+            MaidrKey.TYPE: self.type,
+            MaidrKey.TITLE: 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/mplfinance_lineplot.py

@@ -0,0 +1,273 @@
+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)
+        self._line_titles = []  # Store line titles separately
+
+    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_axes_data(self) -> dict:
+        """
+        Extract axis labels for the plot.
+
+        Returns
+        -------
+        dict
+            Dictionary containing x and y axis labels with custom y-label for moving averages.
+        """
+        x_labels = self.ax.get_xlabel()
+        if not x_labels:
+            x_labels = self.extract_shared_xlabel(self.ax)
+        if not x_labels:
+            x_labels = "Date"
+
+        # Get the period from the first line for y-axis label
+        ma_period = self._extract_moving_average_period()
+        y_label = (
+            f"{ma_period}-day mav price ($)"
+            if ma_period
+            else "Moving Average Price ($)"
+        )
+
+        return {MaidrKey.X: x_labels, MaidrKey.Y: y_label}
+
+    def _extract_moving_average_periods(self) -> List[str]:
+        """
+        Extract all moving average periods from the _maidr_ma_period attributes set by the mplfinance patch.
+
+        Returns
+        -------
+        List[str]
+            List of moving average periods (e.g., ["3", "6", "30"]).
+        """
+        all_lines = self.ax.get_lines()
+        periods = []
+        for line in all_lines:
+            # Get the period that was stored by the mplfinance patch
+            ma_period = getattr(line, "_maidr_ma_period", None)
+            if ma_period is not None:
+                periods.append(str(ma_period))
+
+        # Remove duplicates and sort
+        periods = sorted(list(set(periods)))
+
+        return periods
+
+    def _extract_moving_average_period(self) -> str:
+        """
+        Extract the moving average period from the _maidr_ma_period attribute set by the mplfinance patch.
+
+        Returns
+        -------
+        str
+            The moving average period (e.g., "3", "6", "30") or empty string if no period found.
+        """
+        periods = self._extract_moving_average_periods()
+        return periods[0] if periods else ""
+
+    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
+
+            # Get the period for this specific line
+            ma_period = getattr(line, "_maidr_ma_period", None)
+
+            # Create title for this line
+            line_title = (
+                f"{ma_period}-Day Moving Average Line Plot"
+                if ma_period
+                else "Moving Average Line Plot"
+            )
+
+            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),
+                }
+                line_data.append(point_data)
+
+            if line_data:
+                # Create line data with title, axes, and points structure
+                line_with_metadata = {
+                    "title": line_title,
+                    "axes": {
+                        "x": "Date",
+                        "y": f"{ma_period}-day mav price ($)"
+                        if ma_period
+                        else "Moving Average Price ($)",
+                    },
+                    "points": line_data,
+                }
+                all_lines_data.append(line_with_metadata)
+
+        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)
+
+    def _extract_line_titles(self) -> List[str]:
+        """
+        Extract titles for all moving average lines.
+
+        Returns
+        -------
+        List[str]
+            List of titles for each line.
+        """
+        all_lines = self.ax.get_lines()
+        titles = []
+
+        for line in all_lines:
+            ma_period = getattr(line, "_maidr_ma_period", None)
+            title = (
+                f"{ma_period}-Day Moving Average Line Plot"
+                if ma_period
+                else "Moving Average Line Plot"
+            )
+            titles.append(title)
+
+        return titles
+
+    def render(self) -> dict:
+        """Initialize the MAIDR schema dictionary with basic plot information."""
+        # Use the first line's period for the main title
+        ma_period = self._extract_moving_average_period()
+        title = (
+            f"{ma_period}-Day Moving Averages Line Plot"
+            if ma_period
+            else "Moving Averages Line Plot"
+        )
+
+        maidr_schema = {
+            MaidrKey.TYPE: self.type,
+            MaidrKey.TITLE: 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/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,
+)

maidr/patch/mplfinance.py

@@ -0,0 +1,215 @@
+import uuid
+import wrapt
+import mplfinance as mpf
+import numpy as np
+from matplotlib.collections import LineCollection, PolyCollection
+from matplotlib.patches import Rectangle
+from matplotlib.lines import Line2D
+from maidr.core.enum import PlotType
+from maidr.patch.common import common
+from maidr.core.context_manager import ContextManager
+
+
+def mplfinance_plot_patch(wrapped, instance, args, kwargs):
+    """
+    Enhanced patch function for `mplfinance.plot` that registers separate layers:
+    - CANDLESTICK: For OHLC data (candle bodies and wicks)
+    - BAR: For volume data (volume bars)
+    - LINE: For moving averages (lines)
+
+    This function intercepts calls to `mplfinance.plot`, identifies the resulting
+    candlestick, volume, and moving average components, and registers them with
+    maidr using the common patching mechanism.
+    """
+    # Ensure `returnfig=True` to capture the figure and axes objects.
+    original_returnfig = kwargs.get("returnfig", False)
+    kwargs["returnfig"] = True
+
+    with ContextManager.set_internal_context():
+        result = wrapped(*args, **kwargs)
+
+    # Validate that we received the expected figure and axes tuple
+    if not (isinstance(result, tuple) and len(result) >= 2):
+        return result if original_returnfig else None
+
+    fig, axes = result[0], result[1]
+    ax_list = axes if isinstance(axes, list) else [axes]
+
+    # Enhanced axis identification using content-based detection
+    price_ax = None
+    volume_ax = None
+
+    # Identify axes by their content rather than just labels
+    for ax in ax_list:
+        # Price axis has candlestick collections (LineCollection for wicks, PolyCollection for bodies)
+        if any(isinstance(c, (LineCollection, PolyCollection)) for c in ax.collections):
+            price_ax = ax
+        # Volume axis has rectangle patches for volume bars
+        elif any(isinstance(p, Rectangle) for p in ax.patches):
+            volume_ax = ax
+        # Fallback: use y-label if content-based detection fails
+        elif price_ax is None and "price" in ax.get_ylabel().lower():
+            price_ax = ax
+        elif volume_ax is None and "volume" in ax.get_ylabel().lower():
+            volume_ax = ax
+
+    # Try to extract date numbers from the data
+    date_nums = None
+    data = None
+    if len(args) > 0:
+        data = args[0]
+    elif "data" in kwargs:
+        data = kwargs["data"]
+
+    if data is not None:
+        if hasattr(data, "Date_num"):
+            date_nums = list(data["Date_num"])
+        elif hasattr(data, "index"):
+            # fallback: use index if it's a DatetimeIndex
+            try:
+                import matplotlib.dates as mdates
+
+                date_nums = [mdates.date2num(d) for d in data.index]
+            except Exception:
+                pass
+
+    # Process and register the Candlestick plot
+    if price_ax:
+        wick_collection = next(
+            (c for c in price_ax.collections if isinstance(c, LineCollection)), None
+        )
+        body_collection = next(
+            (c for c in price_ax.collections if isinstance(c, PolyCollection)), None
+        )
+
+        if wick_collection and body_collection:
+            gid = f"maidr-{uuid.uuid4()}"
+            wick_collection.set_gid(gid)
+            body_collection.set_gid(gid)
+
+            candlestick_kwargs = dict(
+                kwargs,
+                _maidr_wick_collection=wick_collection,
+                _maidr_body_collection=body_collection,
+                _maidr_date_nums=date_nums,
+            )
+            common(
+                PlotType.CANDLESTICK,
+                lambda *a, **k: price_ax,
+                instance,
+                args,
+                candlestick_kwargs,
+            )
+
+    # Process and register the Volume plot
+    if volume_ax:
+        volume_patches = [p for p in volume_ax.patches if isinstance(p, Rectangle)]
+
+        if not volume_patches:
+            # Search in shared axes for volume patches
+            for twin_ax in volume_ax.get_shared_x_axes().get_siblings(volume_ax):
+                if twin_ax is not volume_ax:
+                    volume_patches.extend(
+                        [p for p in twin_ax.patches if isinstance(p, Rectangle)]
+                    )
+
+        if volume_patches:
+            # Set GID for volume patches for highlighting
+            for patch in volume_patches:
+                if patch.get_gid() is None:
+                    gid = f"maidr-{uuid.uuid4()}"
+                    patch.set_gid(gid)
+
+            bar_kwargs = dict(
+                kwargs,
+                _maidr_patches=volume_patches,
+                _maidr_date_nums=date_nums,
+            )
+            common(PlotType.BAR, lambda *a, **k: volume_ax, instance, args, bar_kwargs)
+
+    # Process and register Moving Averages as LINE plots
+    if price_ax:
+        # Find moving average lines (Line2D objects)
+        ma_lines = [line for line in price_ax.get_lines() if isinstance(line, Line2D)]
+
+        # Track processed lines to avoid duplicates
+        processed_lines = set()
+        valid_lines = []
+
+        for line in ma_lines:
+            # Try to identify the moving average period based on NaN count
+            xydata = line.get_xydata()
+
+            if xydata is not None:
+                xydata_array = np.asarray(xydata)
+                nan_count = np.sum(
+                    np.isnan(xydata_array[:, 1])
+                )  # Count NaN in y-values
+
+                # Map NaN count to likely moving average period
+                estimated_period = nan_count + 1
+
+                # Store the period directly on the line for easy access
+                setattr(line, "_maidr_ma_period", estimated_period)
+
+                # Create a better label for the line
+                label = str(line.get_label())
+                if label.startswith("_child"):
+                    new_label = f"Moving Average {estimated_period} days"
+                    line.set_label(new_label)
+                else:
+                    # If it's not a _child label, still add the period info
+                    new_label = f"{label}_MA{estimated_period}"
+                    line.set_label(new_label)
+
+            # Create a unique identifier for this line based on its data
+            if xydata is not None:
+                xydata_array = np.asarray(xydata)
+                if xydata_array.size > 0:
+                    # Use shape and first few values to create a unique identifier
+                    first_values = (
+                        xydata_array[:3].flatten()
+                        if xydata_array.size >= 6
+                        else xydata_array.flatten()
+                    )
+                    data_hash = hash(f"{xydata_array.shape}_{str(first_values)}")
+                    line_id = f"{line.get_label()}_{data_hash}"
+                else:
+                    line_id = f"{line.get_label()}"
+            else:
+                line_id = f"{line.get_label()}"
+
+            if line_id in processed_lines:
+                continue
+
+            processed_lines.add(line_id)
+
+            # Validate that the line has valid data
+            if xydata is None or xydata_array.size == 0:
+                continue
+
+            # Store date numbers on the line for the line plot class to use
+            if date_nums is not None:
+                setattr(line, "_maidr_date_nums", date_nums)
+
+            # Ensure GID is set for highlighting
+            if line.get_gid() is None:
+                gid = f"maidr-{uuid.uuid4()}"
+                line.set_gid(gid)
+
+            # Add to valid lines list
+            valid_lines.append(line)
+
+        # Register all valid lines as a single LINE plot
+        if valid_lines:
+            line_kwargs = dict(kwargs)
+            common(PlotType.LINE, lambda *a, **k: price_ax, instance, args, line_kwargs)
+
+    if not original_returnfig:
+        return None
+
+    return result
+
+
+# Apply the patch to mplfinance.plot
+wrapt.wrap_function_wrapper(mpf, "plot", mplfinance_plot_patch)

maidr/util/__init__.py

@@ -0,0 +1,3 @@
+from .plot_detection import PlotDetectionUtils
+
+__all__ = ["PlotDetectionUtils"]