fix: legend visibility, style persistence, and axis ranges in combined figures

Combined figures (overlay, add_secondary_y) had three issues:

1. Legends disappeared because Plotly Express sets showlegend=False on
   single-trace figures. Now unnamed traces get names derived from the
   source figure's y-axis title, and showlegend is fixed per legendgroup.

2. Colors and styles were lost during animation because frame traces
   carried PX defaults. Now marker, line, opacity and legend properties
   are propagated from fig.data into all animation frame traces.

3. Axis ranges were computed from fig.data only, so frames with different
   data ranges went off-screen during animation. Now global min/max is
   computed across all frames and set explicitly on the layout.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: FBumann

tests/test_figures.py

@@ -616,3 +616,142 @@ def test_secondary_not_modified(self) -> None:
 
         # Secondary traces should still use original yaxis
         assert secondary.data[0].yaxis == original_yaxis
+
+
+class TestLegendVisibility:
+    """Tests that combined figures preserve legend visibility."""
+
+    def test_overlay_single_trace_figures_with_names(self) -> None:
+        """Overlay of named single-trace figures shows legend."""
+        da1 = xr.DataArray([1, 2, 3], dims=["x"], name="a")
+        da2 = xr.DataArray([4, 5, 6], dims=["x"], name="b")
+
+        fig1 = xpx(da1).line()
+        fig1.update_traces(name="Series A")
+        fig2 = xpx(da2).line()
+        fig2.update_traces(name="Series B")
+
+        combined = overlay(fig1, fig2)
+
+        assert combined.data[0].showlegend is True
+        assert combined.data[1].showlegend is True
+
+    def test_overlay_unnamed_traces_get_yaxis_title(self) -> None:
+        """Overlay of unnamed traces derives names from y-axis titles."""
+        da1 = xr.DataArray([1, 2, 3], dims=["x"], name="Temperature")
+        da2 = xr.DataArray([4, 5, 6], dims=["x"], name="Pressure")
+
+        fig1 = xpx(da1).line()
+        fig2 = xpx(da2).line()
+
+        combined = overlay(fig1, fig2)
+
+        # Names derived from y-axis titles (DataArray names)
+        assert combined.data[0].name == "Temperature"
+        assert combined.data[1].name == "Pressure"
+        assert combined.data[0].showlegend is True
+        assert combined.data[1].showlegend is True
+
+    def test_overlay_same_name_disambiguated(self) -> None:
+        """Overlay of figures with same y-axis title gets numeric suffix."""
+        da1 = xr.DataArray([1, 2, 3], dims=["x"], name="value")
+        da2 = xr.DataArray([4, 5, 6], dims=["x"], name="value")
+
+        fig1 = xpx(da1).line()
+        fig2 = xpx(da2).line()
+
+        combined = overlay(fig1, fig2)
+
+        assert combined.data[0].name == "value (1)"
+        assert combined.data[1].name == "value (2)"
+
+    def test_overlay_multi_trace_deduplicates_legend(self) -> None:
+        """Overlay of multi-trace figures deduplicates shared legendgroups."""
+        da = xr.DataArray(
+            np.random.rand(10, 3),
+            dims=["x", "cat"],
+            coords={"cat": ["A", "B", "C"]},
+        )
+        fig1 = xpx(da).area()
+        fig2 = xpx(da).line()
+
+        combined = overlay(fig1, fig2)
+
+        # First occurrence of each legendgroup should show, duplicates hidden
+        from collections import defaultdict
+
+        groups: dict[str, list[bool]] = defaultdict(list)
+        for trace in combined.data:
+            lg = trace.legendgroup
+            groups[lg].append(trace.showlegend is True)
+
+        for lg, flags in groups.items():
+            assert flags.count(True) == 1, f"legendgroup {lg!r} has {flags.count(True)} visible"
+
+    def test_add_secondary_y_single_trace_with_names(self) -> None:
+        """add_secondary_y of named single-trace figures shows legend."""
+        da1 = xr.DataArray([1, 2, 3], dims=["x"], name="temp")
+        da2 = xr.DataArray([100, 200, 300], dims=["x"], name="precip")
+
+        fig1 = xpx(da1).line()
+        fig1.update_traces(name="Temperature")
+        fig2 = xpx(da2).bar()
+        fig2.update_traces(name="Precipitation")
+
+        combined = add_secondary_y(fig1, fig2)
+
+        assert combined.data[0].showlegend is True
+        assert combined.data[1].showlegend is True
+
+    def test_overlay_faceted_legendgroup_dedup(self) -> None:
+        """Faceted overlay keeps only one showlegend=True per legendgroup."""
+        da = xr.DataArray(
+            np.random.rand(10, 2, 2),
+            dims=["x", "cat", "facet"],
+            coords={"cat": ["A", "B"], "facet": ["left", "right"]},
+        )
+        fig1 = xpx(da).area(facet_col="facet")
+        fig2 = xpx(da).line(facet_col="facet")
+
+        combined = overlay(fig1, fig2)
+
+        # Check each legendgroup has at least one showlegend=True
+        from collections import defaultdict
+
+        groups: dict[str, list[bool]] = defaultdict(list)
+        for trace in combined.data:
+            lg = trace.legendgroup or ""
+            if lg:
+                groups[lg].append(trace.showlegend is True)
+
+        for lg, flags in groups.items():
+            assert any(flags), f"legendgroup {lg!r} has no showlegend=True trace"
+
+    def test_overlay_animation_frames_preserve_style(self) -> None:
+        """Animation frame traces keep legend and color from fig.data."""
+        da = xr.DataArray(
+            np.random.rand(10, 3),
+            dims=["x", "time"],
+            coords={"time": [0, 1, 2]},
+            name="Population",
+        )
+        da_smooth = da.rolling(x=3, center=True).mean()
+        da_smooth.name = "Smoothed"
+
+        fig1 = xpx(da).bar(animation_frame="time")
+        fig1.update_traces(marker={"color": "steelblue"})
+        fig2 = xpx(da_smooth).line(animation_frame="time")
+        fig2.update_traces(line={"color": "red"})
+
+        combined = overlay(fig1, fig2)
+
+        for frame in combined.frames:
+            for i, ft in enumerate(frame.data):
+                src = combined.data[i]
+                assert ft.name == src.name
+                assert ft.showlegend == src.showlegend
+                assert ft.legendgroup == src.legendgroup
+            # Bar trace should keep steelblue
+            assert frame.data[0].marker.color == "steelblue"
+            # Line trace should keep red
+            assert frame.data[1].line.color == "red"

xarray_plotly/figures.py

@@ -13,6 +13,171 @@
     import plotly.graph_objects as go
 
 
+def _get_yaxis_title(fig: go.Figure) -> str:
+    """Extract the primary y-axis title text from a figure.
+
+    Args:
+        fig: A Plotly figure.
+
+    Returns:
+        The y-axis title text, or empty string if not set.
+    """
+    try:
+        return fig.layout.yaxis.title.text or ""
+    except AttributeError:
+        return ""
+
+
+def _ensure_legend_visibility(
+    combined: go.Figure,
+    source_figs: list[go.Figure],
+    trace_slices: list[slice],
+) -> None:
+    """Fix legend visibility on a combined figure.
+
+    Handles three problems that arise when combining Plotly Express figures:
+
+    1. **Unnamed traces** — PX sets ``name=""`` on single-trace (no color)
+       figures.  We derive a name from each source figure's y-axis title.
+    2. **Hidden named traces** — PX sets ``showlegend=False`` on single-trace
+       figures.  We ensure at least one trace per ``legendgroup`` (or each
+       ungrouped named trace) has ``showlegend=True``.
+    3. **Duplicate legend entries** — when two source figures share the same
+       ``legendgroup`` names, we deduplicate so only the first trace per
+       group shows in the legend.
+
+    Args:
+        combined: The combined Plotly figure (mutated in place).
+        source_figs: The original source figures, in trace order.
+        trace_slices: Slices into ``combined.data`` for each source figure.
+    """
+    from collections import defaultdict
+
+    # --- Step 1: label unnamed traces from source y-axis titles -----------
+    labels = [_get_yaxis_title(f) for f in source_figs]
+
+    # If all labels are the same, disambiguate
+    unique_labels = {lb for lb in labels if lb}
+    if len(unique_labels) == 1:
+        labels = [f"{labels[0]} ({i + 1})" for i in range(len(labels))]
+
+    for label, sl in zip(labels, trace_slices, strict=False):
+        if not label:
+            continue
+        for trace in combined.data[sl]:
+            if not getattr(trace, "name", None):
+                trace.name = label
+                trace.legendgroup = label
+
+    # --- Step 2 & 3: fix showlegend per legendgroup -----------------------
+    grouped: dict[str, list[Any]] = defaultdict(list)
+    ungrouped: list[Any] = []
+
+    for trace in combined.data:
+        lg = getattr(trace, "legendgroup", None) or ""
+        if lg:
+            grouped[lg].append(trace)
+        else:
+            ungrouped.append(trace)
+
+    for traces in grouped.values():
+        has_visible = False
+        for t in traces:
+            if has_visible:
+                # Deduplicate: only first keeps showlegend
+                t.showlegend = False
+            elif getattr(t, "name", None):
+                t.showlegend = True
+                has_visible = True
+
+    # Ungrouped traces with a name should show in the legend
+    for trace in ungrouped:
+        if getattr(trace, "name", None):
+            trace.showlegend = True
+
+    # --- Step 4: propagate style properties to animation frame traces ------
+    # When Plotly animates, frame trace data overwrites fig.data properties.
+    # PX frame traces carry name="", showlegend=False and default colors,
+    # discarding any styling the user applied via update_traces() before
+    # combining.  Propagate display properties from fig.data into every frame.
+    _STYLE_ATTRS = ("name", "legendgroup", "showlegend", "marker", "line", "opacity")
+    for frame in combined.frames or []:
+        for i, frame_trace in enumerate(frame.data):
+            if i < len(combined.data):
+                src = combined.data[i]
+                for attr in _STYLE_ATTRS:
+                    src_val = getattr(src, attr, None)
+                    if src_val is not None:
+                        setattr(frame_trace, attr, src_val)
+
+
+def _fix_animation_axis_ranges(fig: go.Figure) -> None:
+    """Set axis ranges to encompass data across all animation frames.
+
+    Plotly.js computes autorange from ``fig.data`` only and does not
+    recalculate during animation.  When different frames have very different
+    data ranges (e.g. population of Brazil vs China), values can go off-screen.
+    This function computes the global min/max for each axis across all frames
+    and sets explicit ranges on the layout.
+
+    Only numeric axes are handled; categorical/date axes are left to autorange.
+
+    Args:
+        fig: A Plotly figure with animation frames (mutated in place).
+    """
+    import numpy as np
+
+    if not fig.frames:
+        return
+
+    from collections import defaultdict
+
+    # Collect numeric y-values per axis across all traces (fig.data + frames)
+    y_by_axis: dict[str, list[float]] = defaultdict(list)
+    x_by_axis: dict[str, list[float]] = defaultdict(list)
+
+    for trace in _iter_all_traces(fig):
+        yaxis = getattr(trace, "yaxis", None) or "y"
+        xaxis = getattr(trace, "xaxis", None) or "x"
+
+        y = getattr(trace, "y", None)
+        if y is not None:
+            try:
+                arr = np.asarray(y, dtype=float)
+                finite = arr[np.isfinite(arr)]
+                if len(finite):
+                    y_by_axis[yaxis].extend(finite.tolist())
+            except (ValueError, TypeError):
+                pass  # Non-numeric (categorical) — skip
+
+        x = getattr(trace, "x", None)
+        if x is not None:
+            try:
+                arr = np.asarray(x, dtype=float)
+                finite = arr[np.isfinite(arr)]
+                if len(finite):
+                    x_by_axis[xaxis].extend(finite.tolist())
+            except (ValueError, TypeError):
+                pass
+
+    # Apply ranges to layout
+    for axis_ref, values in y_by_axis.items():
+        if not values:
+            continue
+        lo, hi = min(values), max(values)
+        pad = (hi - lo) * 0.05 or 1  # 5% padding
+        layout_prop = "yaxis" if axis_ref == "y" else f"yaxis{axis_ref[1:]}"
+        fig.layout[layout_prop].range = [lo - pad, hi + pad]
+
+    for axis_ref, values in x_by_axis.items():
+        if not values:
+            continue
+        lo, hi = min(values), max(values)
+        pad = (hi - lo) * 0.05 or 1
+        layout_prop = "xaxis" if axis_ref == "x" else f"xaxis{axis_ref[1:]}"
+        fig.layout[layout_prop].range = [lo - pad, hi + pad]
+
+
 def _iter_all_traces(fig: go.Figure) -> Iterator[Any]:
     """Iterate over all traces in a figure, including animation frames.
 
@@ -194,17 +359,11 @@ def overlay(base: go.Figure, *overlays: go.Figure) -> go.Figure:
         _validate_compatible_structure(base, overlay)
         _validate_animation_compatibility(base, overlay)
 
-    # Create new figure with base's layout
-    combined = go.Figure(layout=copy.deepcopy(base.layout))
-
-    # Add all traces from base
-    for trace in base.data:
-        combined.add_trace(copy.deepcopy(trace))
-
-    # Add all traces from overlays
+    # Create new figure with base's layout and all traces
+    all_traces = [copy.deepcopy(t) for t in base.data]
     for overlay in overlays:
-        for trace in overlay.data:
-            combined.add_trace(copy.deepcopy(trace))
+        all_traces.extend(copy.deepcopy(t) for t in overlay.data)
+    combined = go.Figure(data=all_traces, layout=copy.deepcopy(base.layout))
 
     # Handle animation frames
     if base.frames:
@@ -213,6 +372,17 @@ def overlay(base: go.Figure, *overlays: go.Figure) -> go.Figure:
         merged_frames = _merge_frames(base, list(overlays), base_trace_count, overlay_trace_counts)
         combined.frames = merged_frames
 
+    # Build trace slices for legend fix
+    source_figs = [base, *overlays]
+    slices: list[slice] = []
+    offset = 0
+    for fig in source_figs:
+        n = len(fig.data)
+        slices.append(slice(offset, offset + n))
+        offset += n
+
+    _ensure_legend_visibility(combined, source_figs, slices)
+    _fix_animation_axis_ranges(combined)
     return combined
 
 
@@ -315,19 +485,15 @@ def add_secondary_y(
     rightmost_x = max(x_for_y.values(), key=lambda x: int(x[1:]) if x != "x" else 1)
     rightmost_primary_y = next(y for y, x in x_for_y.items() if x == rightmost_x)
 
-    # Create new figure with base's layout
-    combined = go.Figure(layout=copy.deepcopy(base.layout))
-
-    # Add all traces from base (primary y-axis)
-    for trace in base.data:
-        combined.add_trace(copy.deepcopy(trace))
-
-    # Add all traces from secondary, remapped to secondary y-axes
+    # Build all traces: base (primary) + secondary (remapped to secondary y-axes)
+    all_traces = [copy.deepcopy(t) for t in base.data]
     for trace in secondary.data:
         trace_copy = copy.deepcopy(trace)
         original_yaxis = getattr(trace_copy, "yaxis", None) or "y"
         trace_copy.yaxis = y_mapping[original_yaxis]
-        combined.add_trace(trace_copy)
+        all_traces.append(trace_copy)
+
+    combined = go.Figure(data=all_traces, layout=copy.deepcopy(base.layout))
 
     # Get the rightmost secondary y-axis name for linking
     rightmost_secondary_y = y_mapping[rightmost_primary_y]
@@ -368,6 +534,14 @@ def add_secondary_y(
         merged_frames = _merge_secondary_y_frames(base, secondary, y_mapping)
         combined.frames = merged_frames
 
+    base_n = len(base.data)
+    sec_n = len(secondary.data)
+    _ensure_legend_visibility(
+        combined,
+        [base, secondary],
+        [slice(0, base_n), slice(base_n, base_n + sec_n)],
+    )
+    _fix_animation_axis_ranges(combined)
     return combined