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
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import pickle
import numpy as np
import pytest # noqa: F401
import awkward as ak
behavior = {}
@ak.mixin_class(behavior, "vector")
class Vector:
@ak.mixin_class_method(np.add, set())
def add(self, other):
return ak.zip(
{"x": self.x + other.x, "y": self.y + other.y, "z": self.z + other.z},
with_name="vector",
)
def test():
vec = ak.Array(
[
{"x": 1, "y": 2, "z": 3},
{"x": 4, "y": 5, "z": 9},
],
with_name="vector",
behavior=behavior,
)
assert ak.almost_equal(
(vec + vec),
ak.Array(
[
{"x": 2, "y": 4, "z": 6},
{"x": 8, "y": 10, "z": 18},
],
with_name="vector",
behavior=behavior,
),
)
assert ak.almost_equal(pickle.loads(pickle.dumps(vec)), vec)
|
