legend-loc-compass-notation

ImportanceOfBeingErnest · ImportanceOfBeingErnest · commit c2a6c1b6e022 · 2019-02-25T03:16:39.000+01:00
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -154,6 +154,7 @@ def _update_bbox_to_anchor(self, loc_in_canvas):
         'center'          10
         ===============   =============
 
+
 bbox_to_anchor : `.BboxBase`, 2-tuple, or 4-tuple of floats
     Box that is used to position the legend in conjunction with *loc*.
     Defaults to `axes.bbox` (if called as a method to `.Axes.legend`) or
@@ -335,7 +336,9 @@ class Legend(Artist):
              'upper center': 9,
              'center':       10,
              }
-
+    compasscodes = {'nw': 2, 'n': 9, 'ne': 1, 'w': 6, 'c': 10, 'e': 7,
+                    'sw': 3, 's': 8, 'se': 4}
+    allcodes = {**codes, **compasscodes}
     zorder = 5
 
     def __str__(self):
@@ -505,23 +508,23 @@ def __init__(self, parent, handles, labels,
             if not self.isaxes and loc in [0, 'best']:
                 loc = 'upper right'
         if isinstance(loc, str):
-            if loc not in self.codes:
+            if loc.lower() not in self.allcodes:
                 if self.isaxes:
                     cbook.warn_deprecated(
                         "3.1", message="Unrecognized location {!r}. Falling "
                         "back on 'best'; valid locations are\n\t{}\n"
                         "This will raise an exception %(removal)s."
-                        .format(loc, '\n\t'.join(self.codes)))
+                        .format(loc, '\n\t'.join(self.allcodes)))
                     loc = 0
                 else:
                     cbook.warn_deprecated(
                         "3.1", message="Unrecognized location {!r}. Falling "
                         "back on 'upper right'; valid locations are\n\t{}\n'"
                         "This will raise an exception %(removal)s."
-                        .format(loc, '\n\t'.join(self.codes)))
+                        .format(loc, '\n\t'.join(self.allcodes)))
                     loc = 1
             else:
-                loc = self.codes[loc]
+                loc = self.allcodes[loc.lower()]
         if not self.isaxes and loc == 0:
             cbook.warn_deprecated(
                 "3.1", message="Automatic legend placement (loc='best') not "
diff --git a/lib/matplotlib/rcsetup.py b/lib/matplotlib/rcsetup.py
@@ -601,16 +601,16 @@ def validate_markevery(s):
 validate_legend_loc = ValidateInStrings(
     'legend_loc',
     ['best',
-     'upper right',
-     'upper left',
-     'lower left',
-     'lower right',
+     'upper right', 'ne',
+     'upper left', 'nw',
+     'lower left', 'sw',
+     'lower right', 'se',
      'right',
-     'center left',
-     'center right',
-     'lower center',
-     'upper center',
-     'center'], ignorecase=True)
+     'center left', 'w',
+     'center right', 'e',
+     'lower center', 's',
+     'upper center', 'n',
+     'center', 'c'], ignorecase=True)
 
 
 def validate_svg_fonttype(s):
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -571,3 +571,32 @@ def test_no_warn_big_data_when_loc_specified():
         l = ax.legend('best')
         fig.canvas.draw()
     assert len(records) == 0
+
+
+def test_legend_loc_compass_codes():
+    locs = ['NW', 'N', 'NE', 'W', 'C', 'E', 'SW', 'S', 'SE']
+    codes = [2, 9, 1, 6, 10, 7, 3, 8, 4]
+
+    for loc in locs:
+        plt.rcParams["legend.loc"] = loc
+        assert plt.rcParams["legend.loc"] == loc.lower()
+
+    fig1, axes1 = plt.subplots(3, 3)
+    fig2, axes2 = plt.subplots(3, 3)
+
+    locs = ['NW', 'N', 'NE', 'W', 'C', 'E', 'SW', 'S', 'SE']
+    codes = [2, 9, 1, 6, 10, 7, 3, 8, 4]
+
+    for ax, loc in zip(axes1.flat, locs):
+        ax.plot([1, 2], label=loc)
+        ax.legend(loc=loc)
+
+    for ax, loc in zip(axes2.flat, codes):
+        ax.plot([1, 2], label=loc)
+        ax.legend(loc=loc)
+
+    for ax1, ax2 in zip(axes1.flat, axes2.flat):
+        leg1 = ax1.get_legend()
+        leg2 = ax2.get_legend()
+        assert leg1._get_loc() == leg2._get_loc()
+
