-
Notifications
You must be signed in to change notification settings - Fork 78
Expand file tree
/
Copy pathdecorators.py
More file actions
160 lines (128 loc) · 5.24 KB
/
decorators.py
File metadata and controls
160 lines (128 loc) · 5.24 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
import inspect
import os
import re
import string
import sys
import warnings
from functools import partial, wraps
def memoize(f, cache={}):
@wraps(f)
def g(*args, **kwargs):
key = (f, tuple(args), frozenset(kwargs.items()))
if key not in cache:
cache[key] = f(*args, **kwargs)
return cache[key].copy()
return g
class with_doc:
"""
This decorator combines the docstrings of the provided and decorated objects
to produce the final docstring for the decorated object.
"""
def __init__(self, method, use_header=True):
self.method = method
if use_header:
self.header = \
"""
Notes
-----
"""
else:
self.header = ''
def __call__(self, new_method):
new_doc = new_method.__doc__
original_doc = self.method.__doc__
header = self.header
if original_doc and new_doc:
new_method.__doc__ = """
{}
{}
{}
""".format(original_doc, header, new_doc)
elif original_doc:
new_method.__doc__ = original_doc
return new_method
def quantitizer(base_function,
handler_function = lambda *args, **kwargs: 1.0):
"""
wraps a function so that it works properly with physical quantities
(Quantities).
arguments:
base_function - the function to be wrapped
handler_function - a function which takes the same arguments as the
base_function and returns a Quantity (or tuple of Quantities)
which has (have) the units that the output of base_function should
have.
returns:
a wrapped version of base_function that takes the same arguments
and works with physical quantities. It will have almost the same
__name__ and almost the same __doc__.
"""
from .quantity import Quantity
# define a function which will wrap the base function so that it works
# with Quantities
def wrapped_function(*args , **kwargs):
# run the arguments through the handler function, this should
# return a tuple of Quantities which have the correct units
# for the output of the function we are wrapping
handler_quantities= handler_function( *args, **kwargs)
# now we need to turn Quantities into ndarrays so they behave
# correctly
#
# first we simplify all units so that addition and subtraction work
# there may be another way to ensure this, but I do not have any good
# ideas
# in order to modify the args tuple, we have to turn it into a list
args = list(args)
#replace all the quantities in the argument list with ndarrays
for i in range(len(args)):
#test if the argument is a quantity
if isinstance(args[i], Quantity):
#convert the units to the base units
args[i] = args[i].simplified
#view the array as an ndarray
args[i] = args[i].magnitude
#convert the list back to a tuple so it can be used as an output
args = tuple (args)
#replace all the quantities in the keyword argument
#dictionary with ndarrays
for i in kwargs:
#test if the argument is a quantity
if isinstance(kwargs[i], Quantity):
#convert the units to the base units
kwargs[i] = kwargs[i].simplifed()
#view the array as an ndarray
kwargs[i] = kwargs[i].magnitude
#get the result for the function
result = base_function( *args, **kwargs)
# since we have to modify the result, convert it to a list
result = list(result)
#iterate through the handler_quantities and get the correct
# units
length = min( len(handler_quantities) , len(result) )
for i in range(length):
# if the output of the handler is a quantity make the
# output of the wrapper function be a quantity with correct
# units
if isinstance(handler_quantities[i], Quantity):
# the results should have simplified units since that's what
# the inputs were (they were simplified earlier)
# (reasons why this would not be true?)
result[i] = Quantity(
result[i],
handler_quantities[i]
.dimensionality.simplified
)
#now convert the quantity to the appropriate units
result[i] = result[i].rescale(
handler_quantities[i].dimensionality)
#need to convert the result back to a tuple
result = tuple(result)
return result
# give the wrapped function a similar name to the base function
wrapped_function.__name__ = base_function.__name__ + "_QWrap"
# give the wrapped function a similar doc string to the base function's
# doc string but add an annotation to the beginning
wrapped_function.__doc__ = (
"this function has been wrapped to work with Quantities\n"
+ base_function.__doc__)
return wrapped_function