legend-loc-compass-notation

ImportanceOfBeingErnest · ImportanceOfBeingErnest · commit eb232901fb88 · 2018-10-30T22:13:39.000+01:00
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -112,23 +112,24 @@ def _update_bbox_to_anchor(self, loc_in_canvas):
 _legend_kw_doc = '''
 loc : int or string or pair of floats, default: :rc:`legend.loc` ('best' for \
 axes, 'upper right' for figures)
-    The location of the legend. Possible codes are:
-
-        ===============   =============
-        Location String   Location Code
-        ===============   =============
-        'best'            0
-        'upper right'     1
-        'upper left'      2
-        'lower left'      3
-        'lower right'     4
-        'right'           5
-        'center left'     6
-        'center right'    7
-        'lower center'    8
-        'upper center'    9
-        'center'          10
-        ===============   =============
+    The location of the legend.
+    Possible (case-insensitive) strings and codes are:
+
+        ===============  ==============  =============
+        Location String  Compass String  Location Code
+        ===============  ==============  =============
+        'best'                           0
+        'upper right'    'NE'            1
+        'upper left'     'NW'            2
+        'lower left'     'SW'            3
+        'lower right'    'SE'            4
+        'right'                          5
+        'center left'    'W'             6
+        'center right'   'E'             7
+        'lower center'   'S'             8
+        'upper center'   'N'             9
+        'center'         'C'             10
+        ===============  ==============  =============
 
 
     Alternatively can be a 2-tuple giving ``x, y`` of the lower-left
@@ -320,7 +321,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):
@@ -489,23 +492,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
@@ -585,16 +585,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
@@ -563,3 +563,31 @@ def test_alpha_handles():
         lh.set_alpha(1.0)
     assert lh.get_facecolor()[:-1] == hh[1].get_facecolor()[:-1]
     assert lh.get_edgecolor()[:-1] == hh[1].get_edgecolor()[:-1]
+
+
+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()
