matplotlib Zoom - Modify "Zoom to Rectangle" button to always zoom to a square

I am having a really hard time with Matplotlib when using set_aspect('equal') as per this question

I am trying an alternative to this problem. Is there a way for me to modify "Zoom to rectangle" button such that it always zooms to a "square" ? What I meant is, when I use the zoom function and select a rectangle, I want the rectangle automatically change to a squire such that the present aspect ratio is preserved

Can this be done ?

This can be achieved by registering a callback on the axis x or y limit updates and readjusting them to whatever you want. Here is a solution based on another answer but with some additional functionality specific to your use case.

import matplotlib.pyplot as plt

def on_lim_change(*args):
    DESIRED_ASPECT_RATIO = 1.0
    # The axes that were edited are the first argument
    ax: plt.Axes = args[0]
    # Get the x and y limits which will be updated
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()

    # Find the total area encompassed by the selection
    total_area = (xlim[1] - xlim[0]) * (ylim[1] - ylim[0])

    # Calculate the same area but with the correct aspect ratio
    new_xr = (total_area / DESIRED_ASPECT_RATIO) ** 0.5 * DESIRED_ASPECT_RATIO
    new_yr = (total_area / DESIRED_ASPECT_RATIO) ** 0.5

    # Center the new range on the previous center
    xmean = (xlim[0] + xlim[1]) / 2
    ymean = (ylim[0] + ylim[1]) / 2
    # Calculate the extents centered on the previous center
    xlow = xmean - new_xr / 2
    xhigh = xmean + new_xr / 2
    ylow = ymean - new_yr / 2
    yhigh = ymean + new_yr / 2

    # Prevent this callback from calling itself again
    # https://matplotlib.org/stable/api/cbook_api.html#matplotlib.cbook.CallbackRegistry.blocked
    with ax.callbacks.blocked():
        # Set the limits to the new aspect preserved values
        ax.set_xlim(xlow, xhigh)
        ax.set_ylim(ylow, yhigh)

plt.plot([1, 2, 3, 4], [1, -1, 2, -2])
plt.gca().callbacks.connect('xlim_changed', on_lim_change)
plt.gca().callbacks.connect('ylim_changed', on_lim_change)

plt.show()

I played around with this a bit and it works pretty intuitively, but can be a little surprising when you select an area which is very far from the set aspect ratio. You can set the aspect ratio to whatever you want as well, by just adjusting DESIRED_ASPECT_RATIO, in case you have a different situation in the future.