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

maidr/patch/mplfinance.py

@@ -0,0 +1,213 @@
+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
+from maidr.util.mplfinance_utils import MplfinanceDataExtractor
+
+
+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
+
+                # 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"]

maidr/util/mplfinance_utils.py

@@ -0,0 +1,409 @@
+"""
+Utility functions for handling mplfinance-specific data extraction and processing.
+"""
+
+import matplotlib.dates as mdates
+import numpy as np
+from matplotlib.patches import Rectangle
+from typing import List, Optional, Tuple, Any
+
+
+class MplfinanceDataExtractor:
+    """
+    Utility class for extracting and processing mplfinance-specific data.
+
+    This class handles the conversion of mplfinance plot elements (patches, collections)
+    into standardized data formats that can be used by the core plot classes.
+    """
+
+    @staticmethod
+    def extract_volume_data(
+        volume_patches: List[Rectangle], date_nums: Optional[List[float]] = None
+    ) -> List[dict]:
+        """
+        Extract volume data from mplfinance Rectangle patches.
+
+        Parameters
+        ----------
+        volume_patches : List[Rectangle]
+            List of Rectangle patches representing volume bars
+        date_nums : Optional[List[float]], default=None
+            List of matplotlib date numbers corresponding to the patches
+
+        Returns
+        -------
+        List[dict]
+            List of dictionaries with 'x' and 'y' keys for volume data
+        """
+        if not volume_patches:
+            return []
+
+        formatted_data = []
+
+        # Sort patches by x-coordinate to maintain order
+        sorted_patches = sorted(volume_patches, key=lambda p: p.get_x())
+
+        for i, patch in enumerate(sorted_patches):
+            height = patch.get_height()
+
+            # Use date mapping if available, otherwise use index
+            if date_nums is not None and i < len(date_nums):
+                date_num = date_nums[i]
+                x_label = MplfinanceDataExtractor._convert_date_num_to_string(date_num)
+            else:
+                x_label = f"date_{i:03d}"
+
+            formatted_data.append({"x": str(x_label), "y": float(height)})
+
+        return formatted_data
+
+    @staticmethod
+    def extract_candlestick_data(
+        body_collection: Any,
+        wick_collection: Any,
+        date_nums: Optional[List[float]] = None,
+    ) -> List[dict]:
+        """
+        Extract candlestick data from mplfinance collections.
+
+        Parameters
+        ----------
+        body_collection : Any
+            PolyCollection containing candlestick bodies
+        wick_collection : Any
+            LineCollection containing candlestick wicks
+        date_nums : Optional[List[float]], default=None
+            List of matplotlib date numbers corresponding to the candles
+
+        Returns
+        -------
+        List[dict]
+            List of dictionaries with OHLC data
+        """
+        if not body_collection or not hasattr(body_collection, "get_paths"):
+            return []
+
+        candles = []
+        paths = body_collection.get_paths()
+        face_colors = body_collection.get_facecolor()
+
+        for i, path in enumerate(paths):
+            if len(path.vertices) >= 4:
+                # Extract rectangle coordinates from the path
+                vertices = path.vertices
+                x_coords = vertices[:, 0]
+                y_coords = vertices[:, 1]
+
+                x_min, x_max = x_coords.min(), x_coords.max()
+                y_min, y_max = y_coords.min(), y_coords.max()
+
+                # Use date mapping if available
+                if date_nums is not None and i < len(date_nums):
+                    date_num = date_nums[i]
+                    date_str = MplfinanceDataExtractor._convert_date_num_to_string(
+                        date_num
+                    )
+                else:
+                    x_center = (x_min + x_max) / 2
+                    date_str = MplfinanceDataExtractor._convert_date_num_to_string(
+                        x_center
+                    )
+
+                # Determine if this is an up or down candle based on color
+                is_up = MplfinanceDataExtractor._is_up_candle(face_colors, i)
+
+                # Extract OHLC values
+                (
+                    open_val,
+                    close_val,
+                ) = MplfinanceDataExtractor._extract_ohlc_from_rectangle(
+                    y_min, y_max, is_up
+                )
+
+                # Estimate high and low (these would normally come from wick data)
+                high_val = y_max + (y_max - y_min) * 0.1
+                low_val = y_min - (y_max - y_min) * 0.1
+
+                candle_data = {
+                    "value": date_str,
+                    "open": round(open_val, 2),
+                    "high": round(high_val, 2),
+                    "low": round(low_val, 2),
+                    "close": round(close_val, 2),
+                    "volume": 0,  # Volume is handled separately
+                }
+                candles.append(candle_data)
+
+        return candles
+
+    @staticmethod
+    def extract_rectangle_candlestick_data(
+        body_rectangles: List[Rectangle], date_nums: Optional[List[float]] = None
+    ) -> List[dict]:
+        """
+        Extract candlestick data from Rectangle patches (original_flavor).
+
+        Parameters
+        ----------
+        body_rectangles : List[Rectangle]
+            List of Rectangle patches representing candlestick bodies
+        date_nums : Optional[List[float]], default=None
+            List of matplotlib date numbers corresponding to the candles
+
+        Returns
+        -------
+        List[dict]
+            List of dictionaries with OHLC data
+        """
+        if not body_rectangles:
+            return []
+
+        candles = []
+
+        # Sort rectangles by x-coordinate
+        body_rectangles.sort(key=lambda r: r.get_x())
+
+        for i, rect in enumerate(body_rectangles):
+            x_left = rect.get_x()
+            width = rect.get_width()
+            x_center_num = x_left + width / 2.0
+
+            # Convert x coordinate to date
+            if date_nums is not None and i < len(date_nums):
+                date_str = MplfinanceDataExtractor._convert_date_num_to_string(
+                    date_nums[i]
+                )
+            else:
+                date_str = MplfinanceDataExtractor._convert_date_num_to_string(
+                    x_center_num
+                )
+
+            y_bottom = rect.get_y()
+            height = rect.get_height()
+            face_color = rect.get_facecolor()
+
+            # Determine if this is an up or down candle based on color
+            is_up_candle = MplfinanceDataExtractor._is_up_candle_from_color(face_color)
+
+            # Extract OHLC values from rectangle
+            (
+                open_price,
+                close_price,
+            ) = MplfinanceDataExtractor._extract_ohlc_from_rectangle(
+                y_bottom, y_bottom + height, is_up_candle
+            )
+
+            # Estimate high and low
+            high_price = max(open_price, close_price) + height * 0.1
+            low_price = min(open_price, close_price) - height * 0.1
+
+            # Ensure all values are valid numbers
+            open_price = float(open_price) if not np.isnan(open_price) else 0.0
+            high_price = float(high_price) if not np.isnan(high_price) else 0.0
+            low_price = float(low_price) if not np.isnan(low_price) else 0.0
+            close_price = float(close_price) if not np.isnan(close_price) else 0.0
+
+            candle_data = {
+                "value": date_str,
+                "open": round(open_price, 2),
+                "high": round(high_price, 2),
+                "low": round(low_price, 2),
+                "close": round(close_price, 2),
+                "volume": 0,
+            }
+            candles.append(candle_data)
+
+        return candles
+
+    @staticmethod
+    def clean_axis_label(label: str) -> str:
+        """
+        Clean up axis labels by removing LaTeX formatting.
+
+        Parameters
+        ----------
+        label : str
+            The original axis label
+
+        Returns
+        -------
+        str
+            Cleaned axis label
+        """
+        if not label or not isinstance(label, str):
+            return label
+
+        import re
+
+        # Removes LaTeX-like scientific notation, e.g., "$10^{6}$"
+        cleaned_label = re.sub(r"\s*\$.*?\$", "", label).strip()
+        return cleaned_label if cleaned_label else label
+
+    @staticmethod
+    def _convert_date_num_to_string(date_num: float) -> str:
+        """
+        Convert matplotlib date number to date string.
+
+        Parameters
+        ----------
+        date_num : float
+            Matplotlib date number
+
+        Returns
+        -------
+        str
+            Date string in YYYY-MM-DD format or fallback index
+        """
+        try:
+            # Check if this looks like a matplotlib date number (typically > 700000)
+            if date_num > 700000:
+                date_dt = mdates.num2date(date_num)
+                if hasattr(date_dt, "replace"):
+                    date_dt = date_dt.replace(tzinfo=None)
+                return date_dt.strftime("%Y-%m-%d")
+            elif date_num > 1000:
+                # Try converting as if it's a pandas timestamp
+                try:
+                    import pandas as pd
+
+                    date_dt = pd.to_datetime(date_num, unit="D")
+                    return date_dt.strftime("%Y-%m-%d")
+                except:
+                    pass
+        except (ValueError, TypeError, OverflowError):
+            pass
+
+        # Fallback to index-based date string
+        return f"date_{int(date_num):03d}"
+
+    @staticmethod
+    def convert_x_to_date(x_center_num: float, axes: Optional[List] = None) -> str:
+        """
+        Convert matplotlib x-coordinate to date string.
+
+        Parameters
+        ----------
+        x_center_num : float
+            X-coordinate value to convert
+        axes : Optional[List], optional
+            List of matplotlib axes to help with date conversion
+
+        Returns
+        -------
+        str
+            Date string in YYYY-MM-DD format or fallback
+        """
+        # First, try to get the actual dates from the axes x-axis data
+        if axes and len(axes) > 0:
+            ax = axes[0]
+            try:
+                # Get x-axis ticks which might contain the actual dates
+                x_ticks = ax.get_xticks()
+
+                # If we have x-axis ticks and they look like dates (large numbers), use them
+                if len(x_ticks) > 0 and x_ticks[0] > 1000:
+                    # Find the closest tick to our x_center_num
+                    tick_index = int(round(x_center_num))
+                    if 0 <= tick_index < len(x_ticks):
+                        actual_date_num = x_ticks[tick_index]
+
+                        # Convert the actual date number
+                        if actual_date_num > 700000:
+                            date_dt = mdates.num2date(actual_date_num)
+                            if hasattr(date_dt, "replace"):
+                                date_dt = date_dt.replace(tzinfo=None)
+                            date_str = date_dt.strftime("%Y-%m-%d")
+                            return date_str
+            except Exception:
+                pass
+
+        # Use the utility class for date conversion
+        return MplfinanceDataExtractor._convert_date_num_to_string(x_center_num)
+
+    @staticmethod
+    def _is_up_candle(face_colors: Any, index: int) -> bool:
+        """
+        Determine if a candle is up based on face color.
+
+        Parameters
+        ----------
+        face_colors : Any
+            Face colors from the collection
+        index : int
+            Index of the candle
+
+        Returns
+        -------
+        bool
+            True if up candle, False if down candle
+        """
+        is_up = True  # Default to up candle
+        if hasattr(face_colors, "__len__") and len(face_colors) > index:
+            color = (
+                face_colors[index]
+                if hasattr(face_colors[index], "__len__")
+                else face_colors
+            )
+            if isinstance(color, (list, tuple, np.ndarray)):
+                if len(color) >= 3:
+                    # Dark colors typically indicate down candles
+                    if color[0] < 0.5 and color[1] < 0.5 and color[2] < 0.5:
+                        is_up = False
+        return is_up
+
+    @staticmethod
+    def _is_up_candle_from_color(face_color: Any) -> bool:
+        """
+        Determine if a candle is up based on face color (for Rectangle patches).
+
+        Parameters
+        ----------
+        face_color : Any
+            Face color of the rectangle
+
+        Returns
+        -------
+        bool
+            True if up candle, False if down candle
+        """
+        try:
+            if (
+                isinstance(face_color, (list, tuple, np.ndarray))
+                and len(face_color) >= 3
+            ):
+                # Green colors typically indicate up candles
+                if face_color[1] > face_color[0]:
+                    return True
+                else:
+                    return False
+        except (TypeError, IndexError):
+            pass
+        return True  # Default to up candle
+
+    @staticmethod
+    def _extract_ohlc_from_rectangle(
+        y_min: float, y_max: float, is_up: bool
+    ) -> Tuple[float, float]:
+        """
+        Extract open and close values from rectangle bounds.
+
+        Parameters
+        ----------
+        y_min : float
+            Minimum y value of rectangle
+        y_max : float
+            Maximum y value of rectangle
+        is_up : bool
+            Whether this is an up candle
+
+        Returns
+        -------
+        Tuple[float, float]
+            (open_price, close_price)
+        """
+        if is_up:
+            # Up candle: open at bottom, close at top
+            return y_min, y_max
+        else:
+            # Down candle: open at top, close at bottom
+            return y_max, y_min