Merge pull request #6639 from Kojoley/simplify-get_legend_handler

tacaswell · web-flow · commit baf812cc80f7 · 2016-07-12T00:13:09.000-05:00
MNT: refactor get_legend_handler method
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -33,7 +33,8 @@
 
 from matplotlib import rcParams
 from matplotlib.artist import Artist, allow_rasterization
-from matplotlib.cbook import (is_string_like, iterable, silent_list, safezip)
+from matplotlib.cbook import (is_string_like, iterable, silent_list, safezip,
+                              is_hashable)
 from matplotlib.font_manager import FontProperties
 from matplotlib.lines import Line2D
 from matplotlib.patches import Patch, Rectangle, Shadow, FancyBboxPatch
@@ -565,19 +566,19 @@ def get_legend_handler(legend_handler_map, orig_handle):
         method-resolution-order. If no matching key is found, it
         returns None.
         """
-        legend_handler_keys = list(six.iterkeys(legend_handler_map))
-        if orig_handle in legend_handler_keys:
-            handler = legend_handler_map[orig_handle]
-        else:
+        if is_hashable(orig_handle):
+            try:
+                return legend_handler_map[orig_handle]
+            except KeyError:
+                pass
 
-            for handle_type in type(orig_handle).mro():
-                if handle_type in legend_handler_map:
-                    handler = legend_handler_map[handle_type]
-                    break
-            else:
-                handler = None
+        for handle_type in type(orig_handle).mro():
+            try:
+                return legend_handler_map[handle_type]
+            except KeyError:
+                pass
 
-        return handler
+        return None
 
     def _init_legend_box(self, handles, labels, markerfirst=True):
         """
