8000 Proof of concept "accepts units" decorator · matplotlib/matplotlib@f44ee44 · GitHub
[go: up one dir, main page]

Skip to content

Commit f44ee44

Browse files
committed
Proof of concept "accepts units" decorator
Proof of concept "accepts units" decorator Add helper function Mssing commas
1 parent 90437e4 commit f44ee44

File tree

2 files changed

+70
-4
lines changed

2 files changed

+70
-4
lines changed

lib/matplotlib/axes/_axes.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import matplotlib.ticker as mticker
3434
import matplotlib.transforms as mtransforms
3535
import matplotlib.tri as mtri
36+
import matplotlib.units as munits
3637
from matplotlib.cbook import (
3738
mplDeprecation, warn_deprecated, STEP_LOOKUP_MAP, iterable,
3839
safe_first_element)
@@ -3844,6 +3845,7 @@ def dopatch(xs, ys, **kwargs):
38443845
return dict(whiskers=whiskers, caps=caps, boxes=boxes,
38453846
medians=medians, fliers=fliers, means=means)
38463847

3848+
@munits._accepts_units(convert_x=['x'], convert_y=['y'])
38473849
@_preprocess_data(replace_names=["x", "y", "s", "linewidths",
38483850
"edgecolors", "c", "facecolor",
38493851
"facecolors", "color"],
@@ -3990,10 +3992,6 @@ def scatter(self, x, y, s=None, c=None, marker=None, cmap=None, norm=None,
39903992
if edgecolors is None and not rcParams['_internal.classic_mode']:
39913993
edgecolors = 'face'
39923994

3993-
self._process_unit_info(xdata=x, ydata=y, kwargs=kwargs)
3994-
x = self.convert_xunits(x)
3995-
y = self.convert_yunits(y)
3996-
39973995
# np.ma.ravel yields an ndarray, not a masked array,
39983996
# unless its argument is a masked array.
39993997
xy_shape = (np.shape(x), np.shape(y))

lib/matplotlib/units.py

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,12 +43,80 @@ def default_units(x, axis):
4343
"""
4444

4545
from numbers import Number
46+
import inspect
4647

4748
import numpy as np
4849

4950
from matplotlib.cbook import iterable, safe_first_element
5051

5152

53+
def _accepts_units(convert_x, convert_y):
54+
"""
55+
A decorator for functions and methods that accept units. The parameters
56+
indicated in *convert_x* and *convert_y* are used to update the axis
57+
unit information, are converted, and then handed on to the decorated
58+
function.
59+
60+
Parameters
61+
----------
62+
convert_x, convert_y : list
63+
A list of integers or strings, indicating the arguments to be converted
64+
"""
65+
def decorator(func):
66+
def wrapper(*args, **kwargs):
67+
axes = args[0]
68+
# Bind the incoming arguments to the function signature
69+
bound_args = inspect.signature(func).bind(*args, **kwargs)
70+
# Get the original arguments - these will be modified later
71+
arguments = bound_args.arguments
72+
# Check for data kwarg
73+
has_data = (('data' in arguments) and
74+
(arguments['data'] is not None))
75+
if has_data:
76+
data = arguments['data']
77+
78+
# Helper method to process unit info, and convert *original_data*
79+
def _process_info(original_data, axis):
80+
if axis == 'x':
81+
axes._process_unit_info(xdata=original_data, kwargs=kwargs)
82+
converted_data = axes.convert_xunits(original_data)
83+
elif axis == 'y':
84+
axes._process_unit_info(ydata=original_data, kwargs=kwargs)
85+
converted_data = axes.convert_yunits(original_data)
86+
return converted_data
87+
88+
# Loop through each argument to be converted, update the axis
89+
# unit info, convert argument, and replace in *arguments* with
90+
# converted values
91+
for arg in convert_x:
92+
if has_data and arguments[arg] in data:
93+
data_arg = arguments[arg]
94+
data[data_arg] = _process_info(data[data_arg], 'x')
95+
else:
96+
arguments[arg] = _process_info(arguments[arg], 'x')
97+
98+
for arg in convert_y:
99+
if has_data and arguments[arg] in data:
100+
data_arg = arguments[arg]
101+
data[data_arg] = _process_info(data[data_arg], 'y')
102+
else:
103+
arguments[arg] = _process_info(arguments[arg], 'y')
104+
105+
if has_data:
106+
arguments['data'] = data
107+
# Update the arguments with converted values
108+
bound_args.arguments = arguments
109+
110+
# Give updated values to the original function
111+
args = bound_args.args
112+
kwargs = bound_args.kwargs
113+
kwargs.pop('xunits', None)
114+
kwargs.pop('yunits', None)
115+
return func(*args, **kwargs)
116+
return wrapper
117+
return decorator
118+
119+
52120
class AxisInfo(object):
53121
"""
54122
Information to support default axis labeling, tick labeling, and

0 commit comments

Comments
 (0)
0