diff --git a/src/vector/backends/awkward.py b/src/vector/backends/awkward.py index 767e3e0d..ef8dfd7d 100644 --- a/src/vector/backends/awkward.py +++ b/src/vector/backends/awkward.py @@ -584,6 +584,17 @@ def elements(self) -> tuple[ArrayOrRecord]: return (self.tau,) +def _projection_class( + cls: type[VectorProtocol], dimension: int +) -> type[VectorProtocol]: + name = f"ProjectionClass{dimension}D" + projection: type[VectorProtocol] | None = getattr(cls, name, None) + if not projection: + msg = f"{dimension}D conversion for {cls.__name__} is not defined" + raise TypeError(msg) + return projection + + def _class_to_name(cls: type[VectorProtocol]) -> str: # respect the type of classes inheriting VectorAwkward classes is_vector = "vector.backends" in cls.__module__ @@ -742,11 +753,11 @@ def _wrap_result( if any( f in fields for f in ("t", "tau", "M", "m", "mass", "E", "e", "energy") ): - cls = cls.ProjectionClass4D + cls = _projection_class(cls, 4) elif any(f in fields for f in ("z", "pz", "theta", "eta")): - cls = cls.ProjectionClass3D + cls = _projection_class(cls, 3) else: - cls = cls.ProjectionClass2D + cls = _projection_class(cls, 2) return maybe_record( ak.zip( @@ -788,7 +799,7 @@ def _wrap_result( ak.zip( dict(zip(names, arrays, strict=True)), depth_limit=first.layout.purelist_depth, - with_name=_class_to_name(cls.ProjectionClass2D), + with_name=_class_to_name(_projection_class(cls, 2)), behavior=None if vector._awkward_registered else first.behavior, ) ) @@ -835,9 +846,9 @@ def _wrap_result( if any( f in fields for f in ("t", "tau", "M", "m", "mass", "E", "e", "energy") ): - cls = cls.ProjectionClass4D + cls = _projection_class(cls, 4) else: - cls = cls.ProjectionClass3D + cls = _projection_class(cls, 3) return maybe_record( ak.zip( @@ -891,7 +902,7 @@ def _wrap_result( ak.zip( dict(zip(names, arrays, strict=True)), depth_limit=first.layout.purelist_depth, - with_name=_class_to_name(cls.ProjectionClass3D), + with_name=_class_to_name(_projection_class(cls, 3)), behavior=None if vector._awkward_registered else first.behavior, ) ) @@ -947,7 +958,7 @@ def _wrap_result( ak.zip( dict(zip(names, arrays, strict=True)), depth_limit=first.layout.purelist_depth, - with_name=_class_to_name(cls.ProjectionClass4D), + with_name=_class_to_name(_projection_class(cls, 4)), behavior=None if vector._awkward_registered else first.behavior, ) ) diff --git a/tests/test_issues.py b/tests/test_issues.py index 4decaf7a..3059e57b 100644 --- a/tests/test_issues.py +++ b/tests/test_issues.py @@ -313,6 +313,40 @@ def test_issue_704(): assert isinstance(vec_vec.neg3D, vector.backends.awkward.MomentumArray4D) +def test_issue_488(): + ak = pytest.importorskip("awkward") + + coordinates = { + 2: {"x": [1.1], "y": [2.2]}, + 3: {"x": [1.1], "y": [2.2], "z": [3.3]}, + 4: {"x": [1.1], "y": [2.2], "z": [3.3], "t": [4.4]}, + } + + for dimension, fields in coordinates.items(): + mixin = getattr(vector.backends.awkward, f"VectorAwkward{dimension}D") + + class VertexArray(mixin, ak.Array): + pass + + class VertexRecord(mixin, ak.Record): + pass + + # a subclass (like coffea's VertexArray) that defines GenericClass but + # deliberately no ProjectionClass*D, because it has no such interpretation + VertexArray.GenericClass = VertexArray + VertexRecord.GenericClass = VertexRecord + + v = ak.zip( + fields, + with_name="Vertex", + behavior={("*", "Vertex"): VertexArray, "Vertex": VertexRecord}, + ) + with pytest.raises( + TypeError, match=f"{dimension}D conversion for VertexArray is not defined" + ): + v.add(v) + + def test_star_import_without_optional_deps(): """from vector import * must not raise even when sympy/awkward are absent.""" # Block sympy via a find_spec-based meta path finder inserted before vector is imported.