Added square plot parameter to axes

Yunfei Yang · Yunfei Yang · commit d6848b153383 · 2015-04-12T17:21:42.000-04:00
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -1299,13 +1299,15 @@ def axis(self, *v, **kwargs):
             ymin, ymax = self.get_ylim()
             return xmin, xmax, ymin, ymax
 
+        emit = kwargs.get('emit', True)
+
         if len(v) == 1 and is_string_like(v[0]):
             s = v[0].lower()
             if s == 'on':
                 self.set_axis_on()
             elif s == 'off':
                 self.set_axis_off()
-            elif s in ('equal', 'tight', 'scaled', 'normal', 'auto', 'image'):
+            elif s in ('equal', 'tight', 'scaled', 'normal', 'auto', 'image', 'square'):
                 self.set_autoscale_on(True)
                 self.set_aspect('auto')
                 self.autoscale_view(tight=False)
@@ -1322,15 +1324,23 @@ def axis(self, *v, **kwargs):
                     self.autoscale_view(tight=True)
                     self.set_autoscale_on(False)
                     self.set_aspect('equal', adjustable='box', anchor='C')
+                elif s == 'square':
+                    self.set_aspect('equal', adjustable='box', anchor='C')
+                    self.set_autoscale_on(False)
+                    xlim = self.get_xlim()
+                    ylim = self.get_ylim()
+                    minlim = min(xlim[0], ylim[0])
+                    maxlim = max(xlim[1], ylim[1])
+                    self.set_xlim([minlim, maxlim], emit=emit, auto=False)
+                    self.set_ylim([minlim, maxlim], emit=emit, auto=False)
 
             else:
                 raise ValueError('Unrecognized string %s to axis; '
                                  'try on or off' % s)
             xmin, xmax = self.get_xlim()
             ymin, ymax = self.get_ylim()
             return xmin, xmax, ymin, ymax
-
-        emit = kwargs.get('emit', True)
+        
         try:
             v[0]
         except IndexError:
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/square_plot.png b/lib/matplotlib/tests/baseline_images/test_axes/square_plot.png
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3635,6 +3635,13 @@ def test_bar_negative_width():
         assert_equal(b._width, 1)
         assert_equal(b._height, indx + 1)
 
+@image_comparison(baseline_images=['square_plot'], extensions=['png'])
+def test_square_plot():
+    x = np.arange(4)
+    y = np.array([1., 3., 5., 7.])
+    fig, ax = plt.subplots()
+    ax.plot(x, y, 'mo')
+    ax.axis('square')
 
 if __name__ == '__main__':
     import nose
