ENH: let unpack_labeled_data handle var-length *args

jankatins · jankatins · commit b38d518c0406 · 2015-08-03T01:16:43.000+02:00
diff --git a/lib/matplotlib/__init__.py b/lib/matplotlib/__init__.py
@@ -1529,7 +1529,8 @@ def _replacer(data, key):
         return key
 
 
-def unpack_labeled_data(replace_names=None, label_namer="y", positional_parameter_names=None):
+def unpack_labeled_data(replace_names=None, replace_all_args=False, label_namer="y",
+                        positional_parameter_names=None):
     """
     A decorator to add a 'data' kwarg to any a function.  The signature
     of the input function must include the ax argument at the first position ::
@@ -1543,14 +1544,18 @@ def foo(ax, *args, **kwargs)
     replace_names : list of strings, optional, default: None
         The list of parameter names which arguments should be replaced by `data[name]`. If None,
         all arguments are replaced if they are included in `data`.
+    replace_all_args : bool, default: False
+        If True, all arguments in *args get replaced, even if they are not in replace_names.
+        NOTE: this should be used only when the order of the names depends on the number of *args.
     label_namer : string, optional, default: 'y'
         The name of the parameter which argument should be used as label, if label is not set. If
         None, the label keyword argument is not set.
+        NOTE: you MUST pass ``label_namer=None`` if the function can't handle a ``label`` kwarg!
     positional_parameter_names : list of strings, optional, default: None
         The full list of positional parameter names (including the `ax` argument at the first place
         and including all possible positional parameter in `*args`), in the right order. Can also
         include all other keyword parameter. Only needed if the wrapped function does contain
-        `*args` and replace_names is not None.
+        `*args` and (replace_names is not None or replace_all_args is False).
     """
     if replace_names is not None:
         replace_names = set(replace_names)
@@ -1579,7 +1584,8 @@ def param(func):
                 # No argnames should be replaced
                 arg_names = []
             else:
-                assert not (positional_parameter_names is None), "Got replace_names and wrapped function uses *args, need positional_parameter_names!"
+                assert not (positional_parameter_names is None or not replace_all_args), \
+                    "Got replace_names and wrapped function uses *args, need positional_parameter_names!"
                 # remove ax arg
                 arg_names = positional_parameter_names[1:]
 
diff --git a/lib/matplotlib/tests/test_labled_data_unpacking.py b/lib/matplotlib/tests/test_labled_data_unpacking.py
@@ -6,13 +6,14 @@
 
 from .. import unpack_labeled_data
 
+
 # these two get used in multiple tests, so define them here
 @unpack_labeled_data(replace_names=["x","y"])
 def plot_func(ax, x, y, ls="x", label=None, w="xyz"):
     return "x: %s, y: %s, ls: %s, w: %s, label: %s" % (list(x),list(y),ls, w, label)
 
 
-@unpack_labeled_data(replace_names=["x","y"], positional_parameter_names=["ax", "x", "y", "ls", "label", "w"])
+@unpack_labeled_data(replace_names=["x","y"], positional_parameter_names=["x", "y", "ls", "label", "w"])
 def plot_func_varags(ax, *args, **kwargs):
     all_args = [None, None, "x", None, "xyz"]
     for i, v in enumerate(args):
@@ -26,25 +27,6 @@ def plot_func_varags(ax, *args, **kwargs):
 
 all_funcs = [plot_func, plot_func_varags]
 
-d# these two get used in multiple tests, so define them here
-@unpack_labeled_data(replace_names=["x","y"])
-def plot_func(ax, x, y, ls="x", label=None, w="xyz"):
-    return "x: %s, y: %s, ls: %s, w: %s, label: %s" % (list(x),list(y),ls, w, label)
-
-
-@unpack_labeled_data(replace_names=["x","y"], positional_parameter_names=["ax", "x", "y", "ls", "label", "w"])
-def plot_func_varags(ax, *args, **kwargs):
-    all_args = [None, None, "x", None, "xyz"]
-    for i, v in enumerate(args):
-        all_args[i] = v
-    for i, k in enumerate(["x", "y", "ls", "label", "w"]):
-        if k in kwargs:
-            all_args[i] = kwargs[k]
-    x, y, ls, label, w = all_args
-    return "x: %s, y: %s, ls: %s, w: %s, label: %s" % (list(x),list(y),ls, w, label)
-
-
-all_funcs = [plot_func, plot_func_varags]
 
 def test_compiletime_checks():
     """test decorator invocations -> no replacements"""
@@ -112,6 +94,7 @@ def f():
     # This sets a label although the function can't handle it.
     assert_raises(TypeError, f)
 
+
 def test_function_call_without_data():
     """test without data -> no replacements"""
     for func in all_funcs:
@@ -192,13 +175,18 @@ def func_varags_replace_all(ax, *args, **kwargs):
         return "x: %s, y: %s, ls: %s, w: %s, label: %s" % (list(x),list(y),ls, w, label)
 
     # in the first case, we can't get a "y" argument, as we don't know the names of the *args
-    assert_equal(func_varags_replace_all(None, "a","b", w="x", data=data), "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: None")
     assert_equal(func_varags_replace_all(None, x="a",y="b", w="x", data=data) , "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: b")
     assert_equal(func_varags_replace_all(None, "a","b", w="x", label="", data=data) , "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
     assert_equal(func_varags_replace_all(None, "a","b", w="x", label="text", data=data) , "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
     assert_equal(func_varags_replace_all(None, x="a",y="b", w="x", label="", data=data) , "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
     assert_equal(func_varags_replace_all(None, x="a",y="b", w="x", label="text", data=data) , "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
 
+    try:
+        from pandas.util.testing import assert_produces_warning
+        assert_equal(func_varags_replace_all(None, "a","b", w="x", data=data), "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: None")
+    except ImportError:
+        pass
+
 
 def test_no_label_replacements():
     """Test without "label_namer=None" -> no label replacement at all"""
@@ -213,6 +201,7 @@ def func_no_label(ax, x, y, ls="x", label=None, w="xyz"):
     assert_equal(func_no_label(None, "a","b", label="", data=data) , "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
     assert_equal(func_no_label(None, "a","b", label="text", data=data) , "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
 
+
 def test_more_args_than_pos_parameter():
     @unpack_labeled_data(replace_names=["x","y"])
     def func(ax, x, y, z=1):
@@ -222,3 +211,31 @@ def func(ax, x, y, z=1):
     def f():
         func(None, "a","b", "z", "z", data=data)
     assert_raises(RuntimeError, f)
+
+
+def test_function_call_with_replace_all_args():
+    """Test with a "replace_all_args" argument, all *args should be replaced"""
+    data = {"a":[1,2],"b":[8,9], "x":"xyz"}
+
+    def funcy(ax, *args, **kwargs):
+        all_args = [None, None, "x", None, "NOT"]
+        for i, v in enumerate(args):
+            all_args[i] = v
+        for i, k in enumerate(["x", "y", "ls", "label", "w"]):
+            if k in kwargs:
+                all_args[i] = kwargs[k]
+        x, y, ls, label, w = all_args
+        return "x: %s, y: %s, ls: %s, w: %s, label: %s" % (list(x),list(y),ls, w, label)
+
+    func = unpack_labeled_data(replace_all_args=True, replace_names=["w"])(funcy)
+
+    #assert_equal(func(None, "a","b", w="x", data=data), "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: None")
+    assert_equal(func(None, "a","b", w="x", label="", data=data) , "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
+    assert_equal(func(None, "a","b", w="x", label="text", data=data) , "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
+
+    func2 = unpack_labeled_data(replace_all_args=True, replace_names=["w"],
+                               positional_parameter_names=["x", "y", "ls", "label", "w"])(funcy)
+
+    assert_equal(func2(None, "a","b", w="x", data=data), "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: b")
+    assert_equal(func2(None, "a","b", w="x", label="", data=data) , "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
+    assert_equal(func2(None, "a","b", w="x", label="text", data=data) , "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
