1from __future__ import (absolute_import, division, print_function)
2
3import re
4
5import numpy as np
6import pytest
7
8from matplotlib import _preprocess_data
9
10
11# Notes on testing the plotting functions itself
12# *   the individual decorated plotting functions are tested in 'test_axes.py'
13# *   that pyplot functions accept a data kwarg is only tested in
14#     test_axes.test_pie_linewidth_0
15
16
17# these two get used in multiple tests, so define them here
18@_preprocess_data(replace_names=["x", "y"], label_namer="y")
19def plot_func(ax, x, y, ls="x", label=None, w="xyz"):
20    return ("x: %s, y: %s, ls: %s, w: %s, label: %s" % (
21        list(x), list(y), ls, w, label))
22
23
24@_preprocess_data(replace_names=["x", "y"], label_namer="y",
25                  positional_parameter_names=["x", "y", "ls", "label", "w"])
26def plot_func_varargs(ax, *args, **kwargs):
27    all_args = [None, None, "x", None, "xyz"]
28    for i, v in enumerate(args):
29        all_args[i] = v
30    for i, k in enumerate(["x", "y", "ls", "label", "w"]):
31        if k in kwargs:
32            all_args[i] = kwargs[k]
33    x, y, ls, label, w = all_args
34    return ("x: %s, y: %s, ls: %s, w: %s, label: %s" % (
35        list(x), list(y), ls, w, label))
36
37
38all_funcs = [plot_func, plot_func_varargs]
39all_func_ids = ['plot_func', 'plot_func_varargs']
40
41
42def test_compiletime_checks():
43    """test decorator invocations -> no replacements"""
44
45    def func(ax, x, y): pass
46
47    def func_args(ax, x, y, *args): pass
48
49    def func_kwargs(ax, x, y, **kwargs): pass
50
51    def func_no_ax_args(*args, **kwargs): pass
52
53    # this is ok
54    _preprocess_data(replace_names=["x", "y"])(func)
55    _preprocess_data(replace_names=["x", "y"])(func_kwargs)
56    # this has "enough" information to do all the replaces
57    _preprocess_data(replace_names=["x", "y"])(func_args)
58
59    # no positional_parameter_names but needed due to replaces
60    with pytest.raises(AssertionError):
61        # z is unknown
62        _preprocess_data(replace_names=["x", "y", "z"])(func_args)
63
64    with pytest.raises(AssertionError):
65        _preprocess_data(replace_names=["x", "y"])(func_no_ax_args)
66
67    # no replacements at all -> all ok...
68    _preprocess_data(replace_names=[], label_namer=None)(func)
69    _preprocess_data(replace_names=[], label_namer=None)(func_args)
70    _preprocess_data(replace_names=[], label_namer=None)(func_kwargs)
71    _preprocess_data(replace_names=[], label_namer=None)(func_no_ax_args)
72
73    # label namer is unknown
74    with pytest.raises(AssertionError):
75        _preprocess_data(label_namer="z")(func)
76
77    with pytest.raises(AssertionError):
78        _preprocess_data(label_namer="z")(func_args)
79
80    # but "ok-ish", if func has kwargs -> will show up at runtime :-(
81    _preprocess_data(label_namer="z")(func_kwargs)
82    _preprocess_data(label_namer="z")(func_no_ax_args)
83
84
85def test_label_problems_at_runtime():
86    """Tests for behaviour which would actually be nice to get rid of."""
87
88    @_preprocess_data(label_namer="z")
89    def func(*args, **kwargs):
90        pass
91
92    # This is a programming mistake: the parameter which should add the
93    # label is not present in the function call. Unfortunately this was masked
94    # due to the **kwargs usage
95    # This would be nice to handle as a compiletime check (see above...)
96    with pytest.warns(RuntimeWarning):
97        func(None, x="a", y="b")
98
99    def real_func(x, y):
100        pass
101
102    @_preprocess_data(label_namer="x")
103    def func(*args, **kwargs):
104        real_func(**kwargs)
105
106    # This sets a label although the function can't handle it.
107    with pytest.raises(TypeError):
108        func(None, x="a", y="b")
109
110
111@pytest.mark.parametrize('func', all_funcs, ids=all_func_ids)
112def test_function_call_without_data(func):
113    """test without data -> no replacements"""
114    assert (func(None, "x", "y") ==
115            "x: ['x'], y: ['y'], ls: x, w: xyz, label: None")
116    assert (func(None, x="x", y="y") ==
117            "x: ['x'], y: ['y'], ls: x, w: xyz, label: None")
118    assert (func(None, "x", "y", label="") ==
119            "x: ['x'], y: ['y'], ls: x, w: xyz, label: ")
120    assert (func(None, "x", "y", label="text") ==
121            "x: ['x'], y: ['y'], ls: x, w: xyz, label: text")
122    assert (func(None, x="x", y="y", label="") ==
123            "x: ['x'], y: ['y'], ls: x, w: xyz, label: ")
124    assert (func(None, x="x", y="y", label="text") ==
125            "x: ['x'], y: ['y'], ls: x, w: xyz, label: text")
126
127
128@pytest.mark.parametrize('func', all_funcs, ids=all_func_ids)
129def test_function_call_with_dict_data(func):
130    """Test with dict data -> label comes from the value of 'x' parameter """
131    data = {"a": [1, 2], "b": [8, 9], "w": "NOT"}
132    assert (func(None, "a", "b", data=data) ==
133            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: b")
134    assert (func(None, x="a", y="b", data=data) ==
135            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: b")
136    assert (func(None, "a", "b", label="", data=data) ==
137            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
138    assert (func(None, "a", "b", label="text", data=data) ==
139            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
140    assert (func(None, x="a", y="b", label="", data=data) ==
141            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
142    assert (func(None, x="a", y="b", label="text", data=data) ==
143            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
144
145
146@pytest.mark.parametrize('func', all_funcs, ids=all_func_ids)
147def test_function_call_with_dict_data_not_in_data(func):
148    "test for the case that one var is not in data -> half replaces, half kept"
149    data = {"a": [1, 2], "w": "NOT"}
150    assert (func(None, "a", "b", data=data) ==
151            "x: [1, 2], y: ['b'], ls: x, w: xyz, label: b")
152    assert (func(None, x="a", y="b", data=data) ==
153            "x: [1, 2], y: ['b'], ls: x, w: xyz, label: b")
154    assert (func(None, "a", "b", label="", data=data) ==
155            "x: [1, 2], y: ['b'], ls: x, w: xyz, label: ")
156    assert (func(None, "a", "b", label="text", data=data) ==
157            "x: [1, 2], y: ['b'], ls: x, w: xyz, label: text")
158    assert (func(None, x="a", y="b", label="", data=data) ==
159            "x: [1, 2], y: ['b'], ls: x, w: xyz, label: ")
160    assert (func(None, x="a", y="b", label="text", data=data) ==
161            "x: [1, 2], y: ['b'], ls: x, w: xyz, label: text")
162
163
164@pytest.mark.parametrize('func', all_funcs, ids=all_func_ids)
165def test_function_call_with_pandas_data(func, pd):
166    """test with pandas dataframe -> label comes from data["col"].name """
167    data = pd.DataFrame({"a": np.array([1, 2], dtype=np.int32),
168                         "b": np.array([8, 9], dtype=np.int32),
169                         "w": ["NOT", "NOT"]})
170
171    assert (func(None, "a", "b", data=data) ==
172            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: b")
173    assert (func(None, x="a", y="b", data=data) ==
174            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: b")
175    assert (func(None, "a", "b", label="", data=data) ==
176            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
177    assert (func(None, "a", "b", label="text", data=data) ==
178            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
179    assert (func(None, x="a", y="b", label="", data=data) ==
180            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
181    assert (func(None, x="a", y="b", label="text", data=data) ==
182            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
183
184
185def test_function_call_replace_all():
186    """Test without a "replace_names" argument, all vars should be replaced"""
187    data = {"a": [1, 2], "b": [8, 9], "x": "xyz"}
188
189    @_preprocess_data(label_namer="y")
190    def func_replace_all(ax, x, y, ls="x", label=None, w="NOT"):
191        return "x: %s, y: %s, ls: %s, w: %s, label: %s" % (
192            list(x), list(y), ls, w, label)
193
194    assert (func_replace_all(None, "a", "b", w="x", data=data) ==
195            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: b")
196    assert (func_replace_all(None, x="a", y="b", w="x", data=data) ==
197            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: b")
198    assert (func_replace_all(None, "a", "b", w="x", label="", data=data) ==
199            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
200    assert (
201        func_replace_all(None, "a", "b", w="x", label="text", data=data) ==
202        "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
203    assert (
204        func_replace_all(None, x="a", y="b", w="x", label="", data=data) ==
205        "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
206    assert (
207        func_replace_all(None, x="a", y="b", w="x", label="text", data=data) ==
208        "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
209
210    @_preprocess_data(label_namer="y")
211    def func_varags_replace_all(ax, *args, **kwargs):
212        all_args = [None, None, "x", None, "xyz"]
213        for i, v in enumerate(args):
214            all_args[i] = v
215        for i, k in enumerate(["x", "y", "ls", "label", "w"]):
216            if k in kwargs:
217                all_args[i] = kwargs[k]
218        x, y, ls, label, w = all_args
219        return "x: %s, y: %s, ls: %s, w: %s, label: %s" % (
220            list(x), list(y), ls, w, label)
221
222    # in the first case, we can't get a "y" argument,
223    # as we don't know the names of the *args
224    assert (func_varags_replace_all(None, x="a", y="b", w="x", data=data) ==
225            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: b")
226    assert (
227        func_varags_replace_all(None, "a", "b", w="x", label="", data=data) ==
228        "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
229    assert (
230        func_varags_replace_all(None, "a", "b", w="x", label="text",
231                                data=data) ==
232        "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
233    assert (
234        func_varags_replace_all(None, x="a", y="b", w="x", label="",
235                                data=data) ==
236        "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
237    assert (
238        func_varags_replace_all(None, x="a", y="b", w="x", label="text",
239                                data=data) ==
240        "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
241
242    with pytest.warns(RuntimeWarning):
243        assert (func_varags_replace_all(None, "a", "b", w="x", data=data) ==
244                "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: None")
245
246
247def test_no_label_replacements():
248    """Test with "label_namer=None" -> no label replacement at all"""
249
250    @_preprocess_data(replace_names=["x", "y"], label_namer=None)
251    def func_no_label(ax, x, y, ls="x", label=None, w="xyz"):
252        return "x: %s, y: %s, ls: %s, w: %s, label: %s" % (
253            list(x), list(y), ls, w, label)
254
255    data = {"a": [1, 2], "b": [8, 9], "w": "NOT"}
256    assert (func_no_label(None, "a", "b", data=data) ==
257            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: None")
258    assert (func_no_label(None, x="a", y="b", data=data) ==
259            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: None")
260    assert (func_no_label(None, "a", "b", label="", data=data) ==
261            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
262    assert (func_no_label(None, "a", "b", label="text", data=data) ==
263            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
264
265
266def test_more_args_than_pos_parameter():
267    @_preprocess_data(replace_names=["x", "y"], label_namer="y")
268    def func(ax, x, y, z=1):
269        pass
270
271    data = {"a": [1, 2], "b": [8, 9], "w": "NOT"}
272    with pytest.raises(RuntimeError):
273        func(None, "a", "b", "z", "z", data=data)
274
275
276def test_function_call_with_replace_all_args():
277    """Test with a "replace_all_args" argument, all *args should be replaced"""
278    data = {"a": [1, 2], "b": [8, 9], "x": "xyz"}
279
280    def funcy(ax, *args, **kwargs):
281        all_args = [None, None, "x", None, "NOT"]
282        for i, v in enumerate(args):
283            all_args[i] = v
284        for i, k in enumerate(["x", "y", "ls", "label", "w"]):
285            if k in kwargs:
286                all_args[i] = kwargs[k]
287        x, y, ls, label, w = all_args
288        return "x: %s, y: %s, ls: %s, w: %s, label: %s" % (
289            list(x), list(y), ls, w, label)
290
291    func = _preprocess_data(replace_all_args=True, replace_names=["w"],
292                            label_namer="y")(funcy)
293
294    assert (func(None, "a", "b", w="x", label="", data=data) ==
295            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
296    assert (func(None, "a", "b", w="x", label="text", data=data) ==
297            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
298
299    func2 = _preprocess_data(replace_all_args=True, replace_names=["w"],
300                             label_namer="y",
301                             positional_parameter_names=["x", "y", "ls",
302                                                         "label", "w"])(funcy)
303
304    assert (func2(None, "a", "b", w="x", data=data) ==
305            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: b")
306    assert (func2(None, "a", "b", w="x", label="", data=data) ==
307            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: ")
308    assert (func2(None, "a", "b", w="x", label="text", data=data) ==
309            "x: [1, 2], y: [8, 9], ls: x, w: xyz, label: text")
310
311
312def test_docstring_addition():
313    @_preprocess_data()
314    def funcy(ax, *args, **kwargs):
315        """Funcy does nothing"""
316        pass
317
318    assert re.search(r".*All positional and all keyword arguments\.",
319                     funcy.__doc__)
320    assert not re.search(r".*All positional arguments\.", funcy.__doc__)
321    assert not re.search(r".*All arguments with the following names: .*",
322                         funcy.__doc__)
323
324    @_preprocess_data(replace_all_args=True, replace_names=[])
325    def funcy(ax, x, y, z, bar=None):
326        """Funcy does nothing"""
327        pass
328
329    assert re.search(r".*All positional arguments\.",
330                     funcy.__doc__)
331    assert not re.search(r".*All positional and all keyword arguments\.",
332                         funcy.__doc__)
333    assert not re.search(r".*All arguments with the following names: .*",
334                         funcy.__doc__)
335
336    @_preprocess_data(replace_all_args=True, replace_names=["bar"])
337    def funcy(ax, x, y, z, bar=None):
338        """Funcy does nothing"""
339        pass
340
341    assert re.search(r".*All positional arguments\.", funcy.__doc__)
342    assert re.search(r".*All arguments with the following names: 'bar'\.",
343                     funcy.__doc__)
344    assert not re.search(r".*All positional and all keyword arguments\.",
345                         funcy.__doc__)
346
347    @_preprocess_data(replace_names=["x", "bar"])
348    def funcy(ax, x, y, z, bar=None):
349        """Funcy does nothing"""
350        pass
351
352    # lists can print in any order, so test for both x,bar and bar,x
353    assert re.search(r".*All arguments with the following names: '.*', '.*'\.",
354                     funcy.__doc__)
355    assert re.search(r".*'x'.*", funcy.__doc__)
356    assert re.search(r".*'bar'.*", funcy.__doc__)
357    assert not re.search(r".*All positional and all keyword arguments\.",
358                         funcy.__doc__)
359    assert not re.search(r".*All positional arguments\.",
360                         funcy.__doc__)
361
362
363def test_positional_parameter_names_as_function():
364    # Also test the _plot_arg_replacer for plot...
365    from matplotlib.axes._axes import _plot_args_replacer
366
367    @_preprocess_data(replace_names=["x", "y"],
368                      positional_parameter_names=_plot_args_replacer)
369    def funcy(ax, *args, **kwargs):
370        return "{args} | {kwargs}".format(args=args, kwargs=kwargs)
371
372    # the normal case...
373    data = {"x": "X", "hy1": "Y"}
374    assert funcy(None, "x", "hy1", data=data) == "('X', 'Y') | {}"
375    assert funcy(None, "x", "hy1", "c", data=data) == "('X', 'Y', 'c') | {}"
376
377    # no arbitrary long args with data
378    with pytest.raises(ValueError):
379        assert (funcy(None, "x", "y", "c", "x", "y", "x", "y", data=data) ==
380                "('X', 'Y', 'c', 'X', 'Y', 'X', 'Y') | {}")
381
382    # In the two arg case, if a valid color spec is in data, we warn but use
383    # it as data...
384    data = {"x": "X", "y": "Y", "ro": "!!"}
385    with pytest.warns(RuntimeWarning):
386        assert funcy(None, "y", "ro", data=data) == "('Y', '!!') | {}"
387