/
mixins.py
112 lines (88 loc) · 3.83 KB
/
mixins.py
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
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
# v2: replace with src/awkward/_v2/behaviors/mixins.py
from __future__ import absolute_import
import sys
import awkward as ak
def mixin_class(registry, name=None):
"""
Args:
registry (dict): The destination behavior mapping registry. Typically,
this would be the global registry #ak.behavior, but one may wish
to register methods in an alternative way.
name (str): The name to assign to the behaviour class.
This decorator can be used to register a behavior mixin class.
Any inherited behaviors will automatically be made available to the decorated
class.
See the "Mixin decorators" section of #ak.behavior for further details.
"""
def register(cls):
cls_name = cls.__name__
if name is None:
behavior_name = cls_name
else:
behavior_name = name
record = type(
cls_name + "Record",
(cls, ak.highlevel.Record),
{"__module__": cls.__module__},
)
setattr(sys.modules[cls.__module__], cls_name + "Record", record)
registry[behavior_name] = record
array = type(
cls_name + "Array",
(cls, ak.highlevel.Array),
{"__module__": cls.__module__},
)
setattr(sys.modules[cls.__module__], cls_name + "Array", array)
registry["*", behavior_name] = array
for basecls in cls.mro():
for method in basecls.__dict__.values():
if hasattr(method, "_awkward_mixin"):
ufunc, rhs, transpose = method._awkward_mixin
if rhs is None:
registry.setdefault((ufunc, behavior_name), method)
continue
for rhs_name in list(rhs) + [behavior_name]:
registry.setdefault((ufunc, behavior_name, rhs_name), method)
if transpose is not None and rhs_name != behavior_name:
registry.setdefault(
(ufunc, rhs_name, behavior_name), transpose
)
if basecls.__name__ in rhs:
rhs.add(behavior_name)
return cls
return register
def mixin_class_method(ufunc, rhs=None, transpose=True):
"""
Args:
ufunc (numpy.ufunc): A universal function (or NEP18 callable) that is
hooked in Awkward Array, i.e. it can be the first argument of a behavior.
rhs (Set[type] or None): Set of right-hand side argument types, optional
if wrapping a unary function. The left-hand side is expected to
always be `self` of the parent class.
transpose (bool): If true, automatically create a transpose signature
(only makes sense for binary ufuncs).
This decorator can be used to register a mixin class method.
Using this decorator ensures that derived classes that are declared with the
#ak.mixin_class decorator will also have the behaviors that this class has.
"""
def register(method):
if not isinstance(rhs, (set, type(None))):
raise ValueError(
"expected a set of right-hand-side argument types"
+ ak._util.exception_suffix(__file__)
)
if transpose and rhs is not None:
def transposed(left, right):
return method(right, left)
# make a copy of rhs, we will edit it later
method._awkward_mixin = (ufunc, set(rhs), transposed)
else:
method._awkward_mixin = (ufunc, rhs, None)
return method
return register
__all__ = [
x for x in list(globals()) if not x.startswith("_") and x not in ("sys", "ak")
]
def __dir__():
return __all__