From 6afedf374a3ff1930a52105c0aae58e032dc53c7 Mon Sep 17 00:00:00 2001 From: Saransh Chopra Date: Wed, 24 Jan 2024 22:18:18 +0100 Subject: [PATCH] don't use awkward stuff in _methods.py --- src/vector/_methods.py | 101 ++++++++++++++++------------------- tests/backends/test_numpy.py | 2 +- tests/test_methods.py | 4 ++ 3 files changed, 52 insertions(+), 55 deletions(-) diff --git a/src/vector/_methods.py b/src/vector/_methods.py index 5d25abf4..e6e38cf2 100644 --- a/src/vector/_methods.py +++ b/src/vector/_methods.py @@ -4130,60 +4130,53 @@ def _handler_of(*objects: VectorProtocol) -> VectorProtocol: assert handler is not None - # if there is a 2D vector in objects - if _check_instance(any, objects, Vector2D): - # if all the objects are not from the same backend - # choose the 2D object of the backend with highest priority if it exists - # or demote the first encountered object of the backend with highest priority to 2D - if ( - not _check_instance(all, objects, vector.VectorObject) - and not _check_instance(all, objects, vector.VectorNumpy) - and not _check_instance(all, objects, vector.VectorAwkward) - ): - new_type = type(handler.to_Vector2D()) - flag = 0 - # if there is a 2D object of the backend with highest priority - # make it the new handler - for obj in objects: - if type(obj) == new_type: - handler = obj - flag = 1 - # else, demote the dimension of the object of the backend with highest priority - if flag == 0: - handler = handler.to_Vector2D() - # if all objects are from the same backend - # use the 2D one as the handler - else: - for obj in objects: - if isinstance(obj, Vector2D): - handler = obj - # if there is no 2D vector but a 3D vector in objects - elif _check_instance(any, objects, Vector3D): - # if all the objects are not from the same backend - # choose the 3D object of the backend with highest priority if it exists - # or demote the first encountered object of the backend with highest priority to 3D - if ( - not _check_instance(all, objects, vector.VectorObject) - and not _check_instance(all, objects, vector.VectorNumpy) - and not _check_instance(all, objects, vector.VectorAwkward) - ): - new_type = type(handler.to_Vector3D()) - flag = 0 - # if there is a 3D object of the backend with highest priority - # make it the new handler - for obj in objects: - if type(obj) == new_type: - handler = obj - flag = 1 - # else, demote the dimension of the object of the backend with highest priority - if flag == 0: - handler = handler.to_Vector3D() - # if all objects are from the same backend - # use the 3D one as the handler - else: - for obj in objects: - if isinstance(obj, Vector3D): - handler = obj + if _check_instance(all, objects, Vector): + # if there is a 2D vector in objects + if _check_instance(any, objects, Vector2D): + # if all the objects are not from the same backend + # choose the 2D object of the backend with highest priority if it exists + # or demote the first encountered object of the backend with highest priority to 2D + if len({_handler_priority.index(obj.__module__) for obj in objects}) != 1: + new_type = type(handler.to_Vector2D()) + flag = 0 + # if there is a 2D object of the backend with highest priority + # make it the new handler + for obj in objects: + if type(obj) == new_type: + handler = obj + flag = 1 + # else, demote the dimension of the object of the backend with highest priority + if flag == 0: + handler = handler.to_Vector2D() + # if all objects are from the same backend + # use the 2D one as the handler + else: + for obj in objects: + if isinstance(obj, Vector2D): + handler = obj + # if there is no 2D vector but a 3D vector in objects + elif _check_instance(any, objects, Vector3D): + # if all the objects are not from the same backend + # choose the 3D object of the backend with highest priority if it exists + # or demote the first encountered object of the backend with highest priority to 3D + if len({_handler_priority.index(obj.__module__) for obj in objects}) != 1: + new_type = type(handler.to_Vector3D()) + flag = 0 + # if there is a 3D object of the backend with highest priority + # make it the new handler + for obj in objects: + if type(obj) == new_type: + handler = obj + flag = 1 + # else, demote the dimension of the object of the backend with highest priority + if flag == 0: + handler = handler.to_Vector3D() + # if all objects are from the same backend + # use the 3D one as the handler + else: + for obj in objects: + if isinstance(obj, Vector3D): + handler = obj return handler diff --git a/tests/backends/test_numpy.py b/tests/backends/test_numpy.py index aa3523b0..e1a389de 100644 --- a/tests/backends/test_numpy.py +++ b/tests/backends/test_numpy.py @@ -589,7 +589,7 @@ def test_demotion(): assert all(v2 + v3 == p_v2_v3) assert all(v3 + v2 == p_v2_v3) - v2 = vector.zip( + v2 = vector.array( { "x": [10.0, 20.0, 30.0], "y": [-10.0, 20.0, 30.0], diff --git a/tests/test_methods.py b/tests/test_methods.py index f8a3c673..f5b5762d 100644 --- a/tests/test_methods.py +++ b/tests/test_methods.py @@ -5,9 +5,13 @@ from __future__ import annotations +import pytest + import vector from vector import VectorObject2D, VectorObject3D, VectorObject4D +awkward = pytest.importorskip("awkward") + def test_handler_of(): object_a = VectorObject4D.from_xyzt(0.0, 0.0, 0.0, 0.0)