Merge pull request #5975 from tacaswell/enh_kwarg_normalize

mdboom · mdboom · commit b17cc2ca8f0e · 2016-02-15T16:27:34.000-05:00
ENH: add kwarg normalization function to cbook
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -2544,6 +2544,114 @@ def safe_first_element(obj):
     return next(iter(obj))
 
 
+def normalize_kwargs(kw, alias_mapping=None, required=(), forbidden=(),
+                     allowed=None):
+    """Helper function to normalize kwarg inputs
+
+    The order they are resolved are:
+
+     1. aliasing
+     2. required
+     3. forbidden
+     4. allowed
+
+    This order means that only the canonical names need appear in
+    `allowed`, `forbidden`, `required`
+
+    Parameters
+    ----------
+
+    alias_mapping, dict, optional
+        A mapping between a canonical name to a list of
+        aliases, in order of precedence from lowest to highest.
+
+        If the canonical value is not in the list it is assumed to have
+        the highest priority.
+
+    required : iterable, optional
+        A tuple of fields that must be in kwargs.
+
+    forbidden : iterable, optional
+        A list of keys which may not be in kwargs
+
+    allowed : tuple, optional
+        A tuple of allowed fields.  If this not None, then raise if
+        `kw` contains any keys not in the union of `required`
+        and `allowed`.  To allow only the required fields pass in
+        ``()`` for `allowed`
+
+    Raises
+    ------
+    TypeError
+        To match what python raises if invalid args/kwargs are passed to
+        a callable.
+
+    """
+    # deal with default value of alias_mapping
+    if alias_mapping is None:
+        alias_mapping = dict()
+
+    # make a local so we can pop
+    kw = dict(kw)
+    # output dictionary
+    ret = dict()
+
+    # hit all alias mappings
+    for canonical, alias_list in six.iteritems(alias_mapping):
+
+        # the alias lists are ordered from lowest to highest priority
+        # so we know to use the last value in this list
+        tmp = []
+        seen = []
+        for a in alias_list:
+            try:
+                tmp.append(kw.pop(a))
+                seen.append(a)
+            except KeyError:
+                pass
+        # if canonical is not in the alias_list assume highest priority
+        if canonical not in alias_list:
+            try:
+                tmp.append(kw.pop(canonical))
+                seen.append(canonical)
+            except KeyError:
+                pass
+        # if we found anything in this set of aliases put it in the return
+        # dict
+        if tmp:
+            ret[canonical] = tmp[-1]
+            if len(tmp) > 1:
+                warnings.warn("Saw kwargs {seen!r} which are all aliases for "
+                              "{canon!r}.  Kept value from {used!r}".format(
+                                  seen=seen, canon=canonical, used=seen[-1]))
+
+    # at this point we know that all keys which are aliased are removed, update
+    # the return dictionary from the cleaned local copy of the input
+    ret.update(kw)
+
+    fail_keys = [k for k in required if k not in ret]
+    if fail_keys:
+        raise TypeError("The required keys {!r} "
+                        "are not in kwargs".format(fail_keys))
+
+    fail_keys = [k for k in forbidden if k in ret]
+    if fail_keys:
+        raise TypeError("The forbidden keys {!r} "
+                        "are in kwargs".format(fail_keys))
+
+    if allowed is not None:
+        allowed_set = set(required) | set(allowed)
+        fail_keys = [k for k in ret if k not in allowed_set]
+        if fail_keys:
+            raise TypeError("kwargs contains {keys!r} which are not in "
+                            "the required {req!r} or "
+                            "allowed {allow!r} keys".format(
+                                keys=fail_keys, req=required,
+                                allow=allowed))
+
+    return ret
+
+
 def get_label(y, default_name):
     try:
         return y.name
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -2,6 +2,7 @@
                         unicode_literals)
 import itertools
 from weakref import ref
+import warnings
 
 from matplotlib.externals import six
 
@@ -309,6 +310,67 @@ def dummy(self):
         pass
 
 
+def _kwarg_norm_helper(inp, expected, kwargs_to_norm, warn_count=0):
+
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        assert expected == cbook.normalize_kwargs(inp, **kwargs_to_norm)
+        assert len(w) == warn_count
+
+
+def _kwarg_norm_fail_helper(inp, kwargs_to_norm):
+    assert_raises(TypeError, cbook.normalize_kwargs, inp, **kwargs_to_norm)
+
+
+def test_normalize_kwargs():
+    fail_mapping = (
+        ({'a': 1}, {'forbidden': ('a')}),
+        ({'a': 1}, {'required': ('b')}),
+        ({'a': 1, 'b': 2}, {'required': ('a'), 'allowed': ()})
+    )
+
+    for inp, kwargs in fail_mapping:
+        yield _kwarg_norm_fail_helper, inp, kwargs
+
+    warn_passing_mapping = (
+        ({'a': 1, 'b': 2}, {'a': 1}, {'alias_mapping': {'a': ['b']}}, 1),
+        ({'a': 1, 'b': 2}, {'a': 1}, {'alias_mapping': {'a': ['b']},
+                                      'allowed': ('a',)}, 1),
+        ({'a': 1, 'b': 2}, {'a': 2}, {'alias_mapping': {'a': ['a', 'b']}}, 1),
+
+        ({'a': 1, 'b': 2, 'c': 3}, {'a': 1, 'c': 3},
+         {'alias_mapping': {'a': ['b']}, 'required': ('a', )}, 1),
+
+    )
+
+    for inp, exp, kwargs, wc in warn_passing_mapping:
+        yield _kwarg_norm_helper, inp, exp, kwargs, wc
+
+    pass_mapping = (
+        ({'a': 1, 'b': 2}, {'a': 1, 'b': 2}, {}),
+        ({'b': 2}, {'a': 2}, {'alias_mapping': {'a': ['a', 'b']}}),
+        ({'b': 2}, {'a': 2}, {'alias_mapping': {'a': ['b']},
+                              'forbidden': ('b', )}),
+
+        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'required': ('a', ),
+                                              'allowed': ('c', )}),
+
+        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'required': ('a', 'c'),
+                                              'allowed': ('c', )}),
+        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'required': ('a', 'c'),
+                                              'allowed': ('a', 'c')}),
+        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'required': ('a', 'c'),
+                                              'allowed': ()}),
+
+        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'required': ('a', 'c')}),
+        ({'a': 1, 'c': 3}, {'a': 1, 'c': 3}, {'allowed': ('a', 'c')}),
+
+    )
+
+    for inp, exp, kwargs in pass_mapping:
+        yield _kwarg_norm_helper, inp, exp, kwargs
+
+
 def test_to_prestep():
     x = np.arange(4)
     y1 = np.arange(4)
