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
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import numpy as np
import pytest
import awkward as ak
def test_only_floats_validation():
class OnlyFloatsError(Exception): ...
class OnlyFloats(ak.Array):
def __awkward_validation__(self) -> None:
def _content_only_floats(layout, **kwargs) -> None:
del kwargs
if layout.is_numpy:
if not np.issubdtype(layout.dtype, np.floating):
raise OnlyFloatsError(
"OnlyFloats arrays can only contain floating-point numbers."
)
_ = ak.transform(_content_only_floats, self, return_value="none")
behavior = {
"OnlyFloats": OnlyFloats,
("*", "OnlyFloats"): OnlyFloats,
}
# this should pass
_ = ak.with_parameter([1.0], "__list__", "OnlyFloats", behavior=behavior)
# this should fail
with pytest.raises(OnlyFloatsError):
_ = ak.with_parameter([1], "__list__", "OnlyFloats", behavior=behavior)
def test_field_validation():
class OnlyAllowedFieldsError(Exception): ...
class OnlyAllowedFields(ak.Array):
_allowed_fields = frozenset({"pt", "eta", "phi", "mass"})
def __awkward_validation__(self) -> None:
fields = set(self.fields)
if self._allowed_fields != fields:
raise OnlyAllowedFieldsError(
f"OnlyAllowedFields arrays can only have fields: {self._allowed_fields}, "
f"but found fields: {fields}"
)
behavior = {
"OnlyAllowedFields": OnlyAllowedFields,
("*", "OnlyAllowedFields"): OnlyAllowedFields,
}
# this should pass
_ = ak.with_parameter(
[{"pt": 1.0, "eta": 1.0, "phi": 1.0, "mass": 1.0}],
"__list__",
"OnlyAllowedFields",
behavior=behavior,
)
# this should fail
with pytest.raises(
OnlyAllowedFieldsError, match="OnlyAllowedFields arrays can only have fields"
):
_ = ak.with_parameter(
[{"pt": 1.0, "eta": 1.0, "phi": 1.0, "mass": 1.0, "energy": 1.0}],
"__list__",
"OnlyAllowedFields",
behavior=behavior,
)
|
