feat: add units implementation

src/awkward/_connect/numpy.py

@@ -23,6 +23,7 @@
 from awkward._typing import Iterator
 from awkward._util import Sentinel
 from awkward.contents.numpyarray import NumpyArray
+from awkward.units import get_unit_registry
 
 # NumPy 1.13.1 introduced NEP13, without which Awkward ufuncs won't work, which
 # would be worse than lacking a feature: it would cause unexpected output.
@@ -211,6 +212,49 @@ def _array_ufunc_signature(ufunc, inputs):
     return tuple(signature)
 
 
+def _array_ufunc_custom_units(ufunc, inputs, kwargs, behavior):
+    registry = get_unit_registry()
+    if registry is None:
+        return None
+
+    # Check if we have units
+    for x in inputs:
+        if isinstance(x, ak.contents.Content) and x.parameter("__units__"):
+            break
+        elif isinstance(x, registry.Quantity):
+            break
+    # Exit now, if not!
+    else:
+        return None
+
+    # Wrap all Awkward Arrays with
+    nextinputs = []
+    for x in inputs:
+        if isinstance(x, ak.contents.Content):
+            nextinputs.append(
+                registry.Quantity(
+                    ak.with_parameter(x, "__units__", None, behavior=behavior),
+                    x.parameter("__units__"),
+                )
+            )
+        else:
+            nextinputs.append(x)
+    out = ufunc(*nextinputs, **kwargs)
+    if not isinstance(out, tuple):
+        out = (out,)
+
+    nextout = []
+    for qty in out:
+        assert isinstance(qty, registry.Quantity)
+        assert isinstance(qty.magnitude, ak.Array)
+        nextout.append(
+            ak.with_parameter(
+                qty.magnitude, "__units__", str(qty.units), highlevel=False
+            )
+        )
+    return tuple(nextout)
+
+
 def array_ufunc(ufunc, method, inputs, kwargs):
     if method != "__call__" or len(inputs) == 0 or "out" in kwargs:
         return NotImplemented
@@ -221,12 +265,35 @@ def array_ufunc(ufunc, method, inputs, kwargs):
     inputs = _array_ufunc_custom_cast(inputs, behavior, backend)
 
     def action(inputs, **ignore):
+        # Do we have any units in the mix? If so, delegate to `pint` to perform
+        # the ufunc dispatch. This will re-enter `array_ufunc`, but without units
+        # NOTE: there's nothing preventing us from handling units for non-NumpyArray
+        #       contents, but for now we restrict ourselves to NumpyArray (in the
+        #       NumpyArray constructor). By running _before_ the custom machinery,
+        #       custom user ufuncs can avoid needing to worry about units
+        out = _array_ufunc_custom_units(ufunc, inputs, kwargs, behavior)
+        if out is not None:
+            return out
+
         signature = _array_ufunc_signature(ufunc, inputs)
         custom = find_ufunc(behavior, signature)
-        # Do we have a custom ufunc (an override of the given ufunc)?
+        # Do we have a custom specific ufunc (an override of the given ufunc)?
         if custom is not None:
             return _array_ufunc_adjust(custom, inputs, kwargs, behavior)
 
+        # Do we have a custom generic ufunc override (a function that accepts _all_ ufuncs)?
+        for x in inputs:
+            if not isinstance(x, ak.contents.Content):
+                continue
+            apply_ufunc = find_ufunc_generic(ufunc, x, behavior)
+            if apply_ufunc is None:
+                continue
+            out = _array_ufunc_adjust_apply(
+                apply_ufunc, ufunc, method, inputs, kwargs, behavior
+            )
+            if out is not None:
+                return out
+
         if ufunc is numpy.matmul:
             raise NotImplementedError(
                 "matrix multiplication (`@` or `np.matmul`) is not yet implemented for Awkward Arrays"
@@ -259,17 +326,6 @@ def action(inputs, **ignore):
 
             return (NumpyArray(result, backend=backend, parameters=parameters),)
 
-        # Do we have a custom generic ufunc override (a function that accepts _all_ ufuncs)?
-        for x in inputs:
-            if isinstance(x, ak.contents.Content):
-                apply_ufunc = find_ufunc_generic(ufunc, x, behavior)
-                if apply_ufunc is not None:
-                    out = _array_ufunc_adjust_apply(
-                        apply_ufunc, ufunc, method, inputs, kwargs, behavior
-                    )
-                    if out is not None:
-                        return out
-
         if all(
             x.parameter("__array__") is not None
             or x.parameter("__record__") is not None
@@ -295,7 +351,7 @@ def action(inputs, **ignore):
 
         return None
 
-    if sum(int(isinstance(x, ak.contents.Content)) for x in inputs) == 1:
+    if sum(int(isinstance(x, ak.contents.Content)) for x in inputs) == 1 and 0:
         where = None
         for i, x in enumerate(inputs):
             if isinstance(x, ak.contents.Content):

src/awkward/contents/content.py

@@ -102,6 +102,13 @@ def _init(self, parameters: dict[str, Any] | None, backend: Backend):
                 )
             )
         else:
+            if not self.is_numpy and parameters.get("__units__") is not None:
+                raise TypeError(
+                    '{} is not allowed to have parameters["__units__"] != None'.format(
+                        type(self).__name__,
+                    )
+                )
+
             if not self.is_list and parameters.get("__array__") in (
                 "string",
                 "bytestring",

src/awkward/units.py

@@ -0,0 +1,59 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+from __future__ import annotations
+
+from threading import RLock
+
+_lock = RLock()
+_unit_registry = None
+_checked_for_pint = False
+
+
+def set_unit_registry(registry):
+    global _unit_registry
+    with _lock:
+        _unit_registry = registry
+
+
+def get_unit_registry():
+    with _lock:
+        _register_if_available()
+        return _unit_registry
+
+
+def _register_if_available():
+    global _checked_for_pint, _unit_registry
+    with _lock:
+        if _checked_for_pint:
+            return
+
+        try:
+            import pint
+        except ModuleNotFoundError:
+            return
+        else:
+            _unit_registry = pint.UnitRegistry()
+        finally:
+            _checked_for_pint = True
+
+
+def register_and_check():
+    """
+    Build `pint` unit registry
+    """
+    try:
+        import pint  # noqa: F401
+
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            """install the 'pint' package with:
+
+        python3 -m pip install pint
+
+    or
+
+        conda install -c conda-forge pint
+    """
+        ) from None
+
+    _register_if_available()