Fix '_wrap_result' for methods called on ak.Record. (#100)

Author: jpivarski

src/vector/_backends/awkward_.py

@@ -332,6 +332,14 @@ def _class_to_name(cls: typing.Type[VectorProtocol]) -> str:
 # the vector class ############################################################
 
 
+def _yes_record(x: ak.Array) -> typing.Optional[typing.Union[float, ak.Record]]:
+    return x[0]
+
+
+def _no_record(x: ak.Array) -> typing.Optional[ak.Array]:
+    return x
+
+
 class VectorAwkward:
     lib: types.ModuleType = numpy
 
@@ -362,7 +370,18 @@ def _wrap_result(
         if returns == [float] or returns == [bool]:
             return result
 
-        elif (
+        if all(not isinstance(x, ak.Array) for x in result):
+            maybe_record = _yes_record
+            result = [
+                ak.Array(x.layout.array[x.layout.at : x.layout.at + 1])
+                if isinstance(x, ak.Record)
+                else ak.Array([x])
+                for x in result
+            ]
+        else:
+            maybe_record = _no_record
+
+        if (
             len(returns) == 1
             and isinstance(returns[0], type)
             and issubclass(returns[0], Azimuthal)
@@ -396,11 +415,13 @@ def _wrap_result(
             else:
                 cls = cls.ProjectionClass2D
 
-            return ak.zip(
-                dict(zip(names, arrays)),
-                depth_limit=first.layout.purelist_depth,
-                with_name=_class_to_name(cls),
-                behavior=None if vector._awkward_registered else first.behavior,
+            return maybe_record(
+                ak.zip(
+                    dict(zip(names, arrays)),
+                    depth_limit=first.layout.purelist_depth,
+                    with_name=_class_to_name(cls),
+                    behavior=None if vector._awkward_registered else first.behavior,
+                )
             )
 
         elif (
@@ -440,11 +461,13 @@ def _wrap_result(
                         names.append(name)
                         arrays.append(self[name])
 
-            return ak.zip(
-                dict(zip(names, arrays)),
-                depth_limit=first.layout.purelist_depth,
-                with_name=_class_to_name(cls.ProjectionClass2D),
-                behavior=None if vector._awkward_registered else first.behavior,
+            return maybe_record(
+                ak.zip(
+                    dict(zip(names, arrays)),
+                    depth_limit=first.layout.purelist_depth,
+                    with_name=_class_to_name(cls.ProjectionClass2D),
+                    behavior=None if vector._awkward_registered else first.behavior,
+                )
             )
 
         elif (
@@ -491,11 +514,13 @@ def _wrap_result(
             else:
                 cls = cls.ProjectionClass3D
 
-            return ak.zip(
-                dict(zip(names, arrays)),
-                depth_limit=first.layout.purelist_depth,
-                with_name=_class_to_name(cls),
-                behavior=None if vector._awkward_registered else first.behavior,
+            return maybe_record(
+                ak.zip(
+                    dict(zip(names, arrays)),
+                    depth_limit=first.layout.purelist_depth,
+                    with_name=_class_to_name(cls),
+                    behavior=None if vector._awkward_registered else first.behavior,
+                )
             )
 
         elif (
@@ -547,11 +572,13 @@ def _wrap_result(
                         names.append(name)
                         arrays.append(self[name])
 
-            return ak.zip(
-                dict(zip(names, arrays)),
-                depth_limit=first.layout.purelist_depth,
-                with_name=_class_to_name(cls.ProjectionClass3D),
-                behavior=None if vector._awkward_registered else first.behavior,
+            return maybe_record(
+                ak.zip(
+                    dict(zip(names, arrays)),
+                    depth_limit=first.layout.purelist_depth,
+                    with_name=_class_to_name(cls.ProjectionClass3D),
+                    behavior=None if vector._awkward_registered else first.behavior,
+                )
             )
 
         elif (
@@ -611,11 +638,13 @@ def _wrap_result(
                         names.append(name)
                         arrays.append(self[name])
 
-            return ak.zip(
-                dict(zip(names, arrays)),
-                depth_limit=first.layout.purelist_depth,
-                with_name=_class_to_name(cls.ProjectionClass4D),
-                behavior=None if vector._awkward_registered else first.behavior,
+            return maybe_record(
+                ak.zip(
+                    dict(zip(names, arrays)),
+                    depth_limit=first.layout.purelist_depth,
+                    with_name=_class_to_name(cls.ProjectionClass4D),
+                    behavior=None if vector._awkward_registered else first.behavior,
+                )
             )
 
         else:

tests/test_issues.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2019-2021, Jonas Eschle, Jim Pivarski, Eduardo Rodrigues, and Henry Schreiner.
+#
+# Distributed under the 3-clause BSD license, see accompanying file LICENSE
+# or https://github.com/scikit-hep/vector for details.
+
+import pytest
+
+import vector
+
+
+def test_issue_99():
+    ak = pytest.importorskip("awkward")
+    vector.register_awkward()
+    vec = ak.Array([{"x": 1.0, "y": 2.0, "z": 3.0}], with_name="Vector3D")
+    assert vec.to_xyz().tolist() == [{"x": 1.0, "y": 2.0, "z": 3.0}]
+    assert vec[0].to_xyz().tolist() == {"x": 1.0, "y": 2.0, "z": 3.0}
+    assert vec[0].to_rhophiz().tolist() == {
+        "rho": 2.23606797749979,
+        "phi": 1.1071487177940904,
+        "z": 3.0,
+    }