figrecipe 0.6.0__py3-none-any.whl → 0.7.4__py3-none-any.whl

figrecipe/_reproducer/_core.py

@@ -0,0 +1,498 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""Core reproduction logic for figure reproduction."""
+
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.axes import Axes
+
+from .._recorder import CallRecord, FigureRecord
+from .._serializer import load_recipe
+
+
+def reproduce(
+    path: Union[str, Path],
+    calls: Optional[List[str]] = None,
+    skip_decorations: bool = False,
+    apply_overrides: bool = True,
+):
+    """Reproduce a figure from a recipe file.
+
+    Parameters
+    ----------
+    path : str or Path
+        Path to .yaml or .png recipe file. If .png is provided,
+        the corresponding .yaml file will be loaded.
+    calls : list of str, optional
+        If provided, only reproduce these specific call IDs.
+    skip_decorations : bool
+        If True, skip decoration calls (labels, legends, etc.).
+    apply_overrides : bool
+        If True (default), apply .overrides.json if it exists.
+        This preserves manual GUI editor changes.
+
+    Returns
+    -------
+    fig : RecordingFigure
+        Reproduced figure (same type as subplots() returns).
+    axes : RecordingAxes or ndarray of RecordingAxes
+        Reproduced axes (single if 1x1, otherwise numpy array).
+
+    Examples
+    --------
+    >>> import figrecipe as ps
+    >>> fig, ax = ps.reproduce("experiment_001.yaml")
+    >>> fig, ax = ps.reproduce("experiment_001.png")  # Also works
+    >>> plt.show()
+    """
+    path = Path(path)
+
+    # Accept both .png and .yaml - find the yaml file
+    if path.suffix.lower() in (".png", ".jpg", ".jpeg", ".pdf", ".svg"):
+        yaml_path = path.with_suffix(".yaml")
+        if not yaml_path.exists():
+            raise FileNotFoundError(
+                f"Recipe file not found: {yaml_path}. "
+                f"Expected .yaml file alongside {path}"
+            )
+        path = yaml_path
+
+    record = load_recipe(path)
+
+    # Check for override file and merge if exists
+    if apply_overrides:
+        overrides_path = path.with_suffix(".overrides.json")
+        if overrides_path.exists():
+            import json
+
+            with open(overrides_path) as f:
+                data = json.load(f)
+
+            # Apply style overrides
+            manual_overrides = data.get("manual_overrides", {})
+            if manual_overrides:
+                # Merge overrides into record style
+                if record.style is None:
+                    record.style = {}
+                record.style.update(manual_overrides)
+
+            # Apply call overrides (kwargs changes from editor)
+            call_overrides = data.get("call_overrides", {})
+            if call_overrides:
+                for ax_key, ax_record in record.axes.items():
+                    for call in ax_record.calls:
+                        if call.id in call_overrides:
+                            # Merge call kwargs overrides
+                            call.kwargs.update(call_overrides[call.id])
+
+    return reproduce_from_record(
+        record,
+        calls=calls,
+        skip_decorations=skip_decorations,
+    )
+
+
+def reproduce_from_record(
+    record: FigureRecord,
+    calls: Optional[List[str]] = None,
+    skip_decorations: bool = False,
+):
+    """Reproduce a figure from a FigureRecord.
+
+    Parameters
+    ----------
+    record : FigureRecord
+        The figure record to reproduce.
+    calls : list of str, optional
+        If provided, only reproduce these specific call IDs.
+    skip_decorations : bool
+        If True, skip decoration calls.
+
+    Returns
+    -------
+    fig : RecordingFigure
+        Reproduced figure (wrapped).
+    axes : RecordingAxes or ndarray of RecordingAxes
+        Reproduced axes (wrapped, numpy array for multi-axes).
+    """
+    from .._recorder import Recorder
+    from .._wrappers import RecordingAxes, RecordingFigure
+
+    # Determine grid size from axes positions
+    max_row = 0
+    max_col = 0
+    for ax_key in record.axes.keys():
+        parts = ax_key.split("_")
+        if len(parts) >= 3:
+            max_row = max(max_row, int(parts[1]))
+            max_col = max(max_col, int(parts[2]))
+
+    nrows = max_row + 1
+    ncols = max_col + 1
+
+    # Create figure
+    fig, mpl_axes = plt.subplots(
+        nrows,
+        ncols,
+        figsize=record.figsize,
+        dpi=record.dpi,
+        constrained_layout=record.constrained_layout,
+    )
+
+    # Apply layout if recorded (skip if constrained_layout is used)
+    if record.layout is not None and not record.constrained_layout:
+        fig.subplots_adjust(**record.layout)
+
+    # Ensure axes is 2D array
+    if nrows == 1 and ncols == 1:
+        axes_2d = np.array([[mpl_axes]])
+    else:
+        axes_2d = np.atleast_2d(mpl_axes)
+        if nrows == 1:
+            axes_2d = axes_2d.reshape(1, -1)
+        elif ncols == 1:
+            axes_2d = axes_2d.reshape(-1, 1)
+
+    # Apply style BEFORE replaying calls (to match original order:
+    # style is applied during subplots(), then user creates plots/decorations)
+    if record.style is not None:
+        from ..styles import apply_style_mm
+
+        for row in range(nrows):
+            for col in range(ncols):
+                apply_style_mm(axes_2d[row, col], record.style)
+
+    # Result cache for resolving references (e.g., clabel needs ContourSet from contour)
+    result_cache: Dict[str, Any] = {}
+
+    # Replay calls on each axes
+    for ax_key, ax_record in record.axes.items():
+        parts = ax_key.split("_")
+        if len(parts) >= 3:
+            row, col = int(parts[1]), int(parts[2])
+        else:
+            row, col = 0, 0
+
+        ax = axes_2d[row, col]
+
+        # Replay plotting calls
+        for call in ax_record.calls:
+            if calls is not None and call.id not in calls:
+                continue
+            result = _replay_call(ax, call, result_cache)
+            if result is not None:
+                result_cache[call.id] = result
+
+        # Replay decorations
+        if not skip_decorations:
+            for call in ax_record.decorations:
+                if calls is not None and call.id not in calls:
+                    continue
+                result = _replay_call(ax, call, result_cache)
+                if result is not None:
+                    result_cache[call.id] = result
+
+    # Finalize tick configuration and special plot types (avoids categorical axis interference)
+    from ..styles._style_applier import finalize_special_plots, finalize_ticks
+
+    for row in range(nrows):
+        for col in range(ncols):
+            finalize_ticks(axes_2d[row, col])
+            finalize_special_plots(axes_2d[row, col], record.style or {})
+
+    # Apply figure-level labels if recorded
+    if record.suptitle is not None:
+        text = record.suptitle.get("text", "")
+        kwargs = record.suptitle.get("kwargs", {}).copy()
+        # Only add y=1.02 if not using constrained_layout (which handles positioning)
+        if "y" not in kwargs and not record.constrained_layout:
+            kwargs["y"] = 1.02
+        fig.suptitle(text, **kwargs)
+
+    if record.supxlabel is not None:
+        text = record.supxlabel.get("text", "")
+        kwargs = record.supxlabel.get("kwargs", {})
+        fig.supxlabel(text, **kwargs)
+
+    if record.supylabel is not None:
+        text = record.supylabel.get("text", "")
+        kwargs = record.supylabel.get("kwargs", {})
+        fig.supylabel(text, **kwargs)
+
+    # Wrap in Recording types (same as subplots() returns)
+    recorder = Recorder()
+    recorder._figure_record = record
+
+    # Wrap axes in RecordingAxes
+    wrapped_axes = np.empty((nrows, ncols), dtype=object)
+    for i in range(nrows):
+        for j in range(ncols):
+            wrapped_axes[i, j] = RecordingAxes(axes_2d[i, j], recorder, position=(i, j))
+
+    # Create RecordingFigure
+    wrapped_fig = RecordingFigure(fig, recorder, wrapped_axes.tolist())
+
+    # Reproduce panel labels if recorded
+    if record.panel_labels is not None:
+        labels = record.panel_labels.get("labels")
+        loc = record.panel_labels.get("loc", "upper left")
+        offset = tuple(record.panel_labels.get("offset", (-0.1, 1.05)))
+        fontsize = record.panel_labels.get("fontsize")
+        fontweight = record.panel_labels.get("fontweight", "bold")
+        color = record.panel_labels.get("color")
+        extra_kwargs = record.panel_labels.get("kwargs", {})
+        if color is not None:
+            extra_kwargs["color"] = color
+        wrapped_fig.add_panel_labels(
+            labels=labels,
+            loc=loc,
+            offset=offset,
+            fontsize=fontsize,
+            fontweight=fontweight,
+            **extra_kwargs,
+        )
+
+    # Return in appropriate format (matching subplots() behavior)
+    if nrows == 1 and ncols == 1:
+        return wrapped_fig, wrapped_axes[0, 0]
+    elif nrows == 1:
+        return wrapped_fig, np.array(wrapped_axes[0], dtype=object)
+    elif ncols == 1:
+        return wrapped_fig, np.array(wrapped_axes[:, 0], dtype=object)
+    else:
+        return wrapped_fig, wrapped_axes
+
+
+def _replay_call(
+    ax: Axes, call: CallRecord, result_cache: Optional[Dict[str, Any]] = None
+) -> Any:
+    """Replay a single call on an axes.
+
+    Parameters
+    ----------
+    ax : Axes
+        The matplotlib axes.
+    call : CallRecord
+        The call to replay.
+    result_cache : dict, optional
+        Cache mapping call_id -> result for resolving references.
+
+    Returns
+    -------
+    Any
+        Result of the matplotlib call.
+    """
+    if result_cache is None:
+        result_cache = {}
+
+    method_name = call.function
+
+    # Check if it's a seaborn call
+    if method_name.startswith("sns."):
+        from ._seaborn import replay_seaborn_call
+
+        return replay_seaborn_call(ax, call)
+
+    # Handle violinplot with inner option specially
+    if method_name == "violinplot":
+        from ._violin import replay_violinplot_call
+
+        return replay_violinplot_call(ax, call)
+
+    # Handle joyplot specially (custom method)
+    if method_name == "joyplot":
+        from ._custom_plots import replay_joyplot_call
+
+        return replay_joyplot_call(ax, call)
+
+    # Handle swarmplot specially (custom method)
+    if method_name == "swarmplot":
+        from ._custom_plots import replay_swarmplot_call
+
+        return replay_swarmplot_call(ax, call)
+
+    method = getattr(ax, method_name, None)
+
+    if method is None:
+        # Method not found, skip
+        return None
+
+    # Reconstruct args
+    args = []
+    for arg_data in call.args:
+        value = _reconstruct_value(arg_data, result_cache)
+        args.append(value)
+
+    # Get kwargs and reconstruct arrays
+    kwargs = _reconstruct_kwargs(call.kwargs)
+
+    # Handle special transform markers
+    if "transform" in kwargs:
+        transform_val = kwargs["transform"]
+        if transform_val == "axes":
+            kwargs["transform"] = ax.transAxes
+        elif transform_val == "data":
+            kwargs["transform"] = ax.transData
+        elif transform_val == "figure":
+            kwargs["transform"] = ax.figure.transFigure
+        # If it's already a Transform object or something else, leave it
+
+    # Call the method
+    try:
+        return method(*args, **kwargs)
+    except Exception as e:
+        # Log warning but continue
+        import warnings
+
+        warnings.warn(f"Failed to replay {method_name}: {e}")
+        return None
+
+
+def _reconstruct_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    """Reconstruct kwargs, converting 2D lists back to numpy arrays.
+
+    Parameters
+    ----------
+    kwargs : dict
+        Raw kwargs from call record.
+
+    Returns
+    -------
+    dict
+        Kwargs with arrays properly reconstructed.
+    """
+    result = {}
+    for key, value in kwargs.items():
+        # Handle 'colors' parameter specially - must be a list for pie/bar/etc.
+        # A single color string like 'red' would be interpreted as ['r','e','d']
+        if key == "colors" and isinstance(value, str):
+            result[key] = [value]
+        elif isinstance(value, list) and len(value) > 0:
+            # Check if it's a 2D list (list of lists) - should be numpy array
+            if isinstance(value[0], list):
+                result[key] = np.array(value)
+            else:
+                # 1D list - could be array or just list, try to preserve
+                result[key] = value
+        else:
+            result[key] = value
+    return result
+
+
+def _reconstruct_value(
+    arg_data: Dict[str, Any], result_cache: Optional[Dict[str, Any]] = None
+) -> Any:
+    """Reconstruct a value from serialized arg data.
+
+    Parameters
+    ----------
+    arg_data : dict
+        Serialized argument data.
+    result_cache : dict, optional
+        Cache mapping call_id -> result for resolving references.
+
+    Returns
+    -------
+    Any
+        Reconstructed value.
+    """
+    if result_cache is None:
+        result_cache = {}
+
+    # Check if we have a pre-loaded array
+    if "_loaded_array" in arg_data:
+        return arg_data["_loaded_array"]
+
+    data = arg_data.get("data")
+
+    # Check if it's a reference to another call's result (e.g., ContourSet for clabel)
+    if isinstance(data, dict) and "__ref__" in data:
+        ref_id = data["__ref__"]
+        if ref_id in result_cache:
+            return result_cache[ref_id]
+        else:
+            import warnings
+
+            warnings.warn(f"Could not resolve reference to {ref_id}")
+            return None
+
+    # Check if it's a list of arrays (e.g., boxplot, violinplot)
+    if arg_data.get("_is_array_list") and isinstance(data, list):
+        dtype = arg_data.get("dtype")
+        # Convert each inner list to numpy array
+        return [
+            np.array(arr_data, dtype=dtype if isinstance(dtype, str) else None)
+            for arr_data in data
+        ]
+
+    # If data is a list, convert to numpy array
+    if isinstance(data, list):
+        dtype = arg_data.get("dtype")
+        try:
+            return np.array(data, dtype=dtype if dtype else None)
+        except (TypeError, ValueError):
+            return np.array(data)
+
+    return data
+
+
+def get_recipe_info(path: Union[str, Path]) -> Dict[str, Any]:
+    """Get information about a recipe without reproducing.
+
+    Parameters
+    ----------
+    path : str or Path
+        Path to .yaml recipe file.
+
+    Returns
+    -------
+    dict
+        Recipe information including:
+        - id: Figure ID
+        - created: Creation timestamp
+        - matplotlib_version: Version used
+        - figsize: Figure size
+        - n_axes: Number of axes
+        - calls: List of call IDs
+    """
+    record = load_recipe(path)
+
+    all_calls = []
+    for ax_record in record.axes.values():
+        for call in ax_record.calls:
+            all_calls.append(
+                {
+                    "id": call.id,
+                    "function": call.function,
+                    "n_args": len(call.args),
+                    "kwargs": list(call.kwargs.keys()),
+                }
+            )
+        for call in ax_record.decorations:
+            all_calls.append(
+                {
+                    "id": call.id,
+                    "function": call.function,
+                    "type": "decoration",
+                }
+            )
+
+    return {
+        "id": record.id,
+        "created": record.created,
+        "matplotlib_version": record.matplotlib_version,
+        "figsize": record.figsize,
+        "dpi": record.dpi,
+        "n_axes": len(record.axes),
+        "calls": all_calls,
+    }
+
+
+__all__ = [
+    "reproduce",
+    "reproduce_from_record",
+    "get_recipe_info",
+]