annotation-like objects axis reference mapping

to a subplot on a new figure. This will allow us to copy the annotations
to a new figure with more axes as part of a px.overlay command.

This just needs to be generalized to work with shapes and images but the
workings can be seen in proto/px_overlay/map_axis_pair_example.py.

Author: nicholas-esterer

proto/px_overlay/map_axis_pair_example.py

@@ -0,0 +1,29 @@
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+import px_overlay
+import pytest
+
+fig0 = px_overlay.make_subplots_all_secondary_y(3, 4)
+fig1 = px_overlay.make_subplots_all_secondary_y(4, 5)
+
+for dims, f in zip([(3, 4), (4, 5)], [fig0, fig1]):
+    for r, c in px_overlay.multi_index(*dims):
+        for sy in [False, True]:
+            f.add_trace(go.Scatter(x=[], y=[]), row=r + 1, col=c + 1, secondary_y=sy)
+
+fig0.add_annotation(row=2, col=3, text="hi", x=0.25, xref="x domain", y=3)
+fig0.add_annotation(
+    row=3, col=4, text="hi", x=0.25, xref="x domain", y=2, secondary_y=True
+)
+
+for an in fig0.layout.annotations:
+    oldaxpair = tuple([an[ref] for ref in ["xref", "yref"]])
+    newaxpair = px_overlay.map_axis_pair(fig0, fig1, oldaxpair)
+    newan = go.layout.Annotation(an)
+    print(oldaxpair)
+    print(newaxpair)
+    newan["xref"], newan["yref"] = newaxpair
+    fig1.add_annotation(newan)
+
+fig0.show()
+fig1.show()

proto/px_overlay/px_overlay.py

@@ -10,6 +10,7 @@
 import json
 from itertools import product, cycle, chain
 from functools import reduce
+import re
 
 
 def multi_index(*kwargs):
@@ -49,6 +50,118 @@ def extract_axis_titles(fig):
     return (r_titles, c_titles)
 
 
+def make_subplots_all_secondary_y(rows, cols):
+    """
+    Get subplots like make_subplots but all also have secondary y-axes.
+    """
+    grid_ref_shape = [rows, cols]
+    specs = [
+        [dict(secondary_y=True) for __ in range(grid_ref_shape[1])]
+        for _ in range(grid_ref_shape[0])
+    ]
+    fig = make_subplots(*grid_ref_shape, specs=specs)
+    return fig
+
+
+def parse_axis_ref(ax):
+    """ Find the axis letter, optional number, and domain of axis. """
+    # TODO: can this be obtained via codegen?
+    pat = re.compile("([xy])(axis)?([0-9]*)( domain)?")
+    matches = pat.match(ax)
+    if matches is None:
+        raise ValueError('Axis "%s" cannot be parsed.' % (ax,))
+    return (matches[1], matches[3], matches[4])
+
+
+def norm_axis_ref(ax):
+    """ normalize ax so it is in the format: yaxis, yaxis2, xaxis7 etc. """
+    al, an, _ = parse_axis_ref(ax)
+    return al + "axis" + an
+
+
+def axis_pair_to_row_col(fig, axpair):
+    """
+    returns the row and column of the subplot having the axis pair and whether it is a
+    secondary y
+    """
+    if "paper" in axpair:
+        raise ValueError('Cannot find row and column of "paper" axis reference.')
+    naxpair = tuple([norm_axis_ref(ax) for ax in axpair])
+    nrows, ncols = fig_grid_ref_shape(fig)
+    row = None
+    col = None
+    for r, c in multi_index(nrows, ncols):
+        for sp in fig._grid_ref[r][c]:
+            if naxpair == sp.layout_keys:
+                row = r + 1
+                col = c + 1
+    if row is None or col is None:
+        raise ValueError("Could not find subplot containing axes (%s,%s)." % nax)
+    secondary_y = False
+    yax = naxpair[1]
+    if fig.layout[yax]["side"] == "right":
+        secondary_y = True
+    return (row, col, secondary_y)
+
+
+def find_subplot_axes(fig, row, col, secondary_y=False):
+    """
+    Returns 2-tuple containing (xaxis,yaxis) at specified row, col and secondary y-axis. 
+    """
+    nrows, ncols = fig_grid_ref_shape(fig)
+    try:
+        sps = fig._grid_ref[row - 1][col - 1]
+    except IndexError:
+        raise IndexError(
+            "Figure does not have a subplot at the requested row or column."
+        )
+
+    def _check_is_secondary_y(sp):
+        xax, yax = sp.layout_keys
+        # TODO: It may not be totally accurate to assume if an y-axis' "side" is
+        # "right" than it is a secondary y axis...
+        return fig.layout[yax]["side"] == "right"
+
+    # find the secondary y axis
+    err_msg = (
+        "Could not find a y-axis " "at the subplot in the requested row or column."
+    )
+    filter_fun = lambda sp: not _check_is_secondary_y(sp)
+    if secondary_y:
+        err_msg = (
+            "Could not find a secondary y-axis "
+            "at the subplot in the requested row or column."
+        )
+        filter_fun = _check_is_secondary_y
+    try:
+        sp = list(filter(filter_fun, sps))[0]
+    except (IndexError, TypeError):
+        # Catch IndexError if the list is empty, catch TypeError if sps isn't
+        # iterable (e.g., is None)
+        raise IndexError(err_msg)
+    return sp.layout_keys
+
+
+def map_axis_pair(old_fig, new_fig, axpair, make_axis_ref=True):
+    """
+    Find the axes on the new figure that will give the same subplot and
+    possibly secondary y axis as on the old figure. This can only
+    work if the axis pair is ("paper","paper") or the axis pair corresponds to a
+    subplot on the old figure the new figure has corresponding rows,
+    columns and secondary y-axes.
+    if make_axis_ref is True, axis is removed from the resulting strings, e.g., xaxis2 -> x2
+    """
+    if axpair == ("paper", "paper"):
+        return ax
+    row, col, secondary_y = axis_pair_to_row_col(old_fig, axpair)
+    newaxpair = find_subplot_axes(new_fig, row, col, secondary_y)
+    axpair_extras = [" domain" if ax.endswith("domain") else "" for ax in axpair]
+    newaxpair = tuple(ax + extra for ax, extra in zip(newaxpair, axpair_extras))
+    if make_axis_ref:
+        newaxpair = tuple(ax.replace("axis", "") for ax in newaxpair)
+    return newaxpair
+
+
 def px_simple_combine(fig0, fig1, fig1_secondary_y=False):
     """
     Combines two figures by just using the layout of the first figure and

proto/px_overlay/test_find_subplot_axes.py

@@ -0,0 +1,40 @@
+from plotly.subplots import make_subplots
+import px_overlay
+import pytest
+
+fig = px_overlay.make_subplots_all_secondary_y(3, 4)
+fig_no_sy = px_overlay.make_subplots(3, 4)
+fig_custom = make_subplots(
+    rows=2,
+    cols=2,
+    specs=[[{}, {}], [{"colspan": 2}, None]],
+    subplot_titles=("First Subplot", "Second Subplot", "Third Subplot"),
+)
+
+
+def test_bad_row_col():
+    with pytest.raises(
+        IndexError,
+        match=r"^Figure does not have a subplot at the requested row or column\.$",
+    ):
+        px_overlay.find_subplot_axes(fig, 4, 2, secondary_y=False)
+    with pytest.raises(
+        IndexError,
+        match=r"^Figure does not have a subplot at the requested row or column\.$",
+    ):
+        px_overlay.find_subplot_axes(fig, 4, 2, secondary_y=True)
+
+
+def test_no_secondary_y():
+    with pytest.raises(
+        IndexError,
+        match=r"^Could not find a secondary y-axis at the subplot in the requested row or column\.$",
+    ):
+        px_overlay.find_subplot_axes(fig_no_sy, 2, 2, secondary_y=True)
+    with pytest.raises(
+        IndexError,
+        match=r"^Could not find a y-axis at the subplot in the requested row or column\.$",
+    ):
+        px_overlay.find_subplot_axes(fig_custom, 2, 2, secondary_y=False)
+    axes = px_overlay.find_subplot_axes(fig_custom, 1, 2, secondary_y=False)
+    assert axes == ("xaxis2", "yaxis2")