@@ -43,12 +43,80 @@ def default_units(x, axis):
4343"""
4444
4545from numbers import Number
46+ import inspect
4647
4748import numpy as np
4849
4950from 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+
52120class AxisInfo (object ):
53121 """
54122 Information to support default axis labeling, tick labeling, and
0 commit comments