refactor: add new module and name arguments to `high_level_functi…

…on` decorator (#2620)

src/awkward/_dispatch.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections.abc import Callable, Collection, Generator
 from functools import wraps
 from inspect import isgenerator
@@ -10,13 +12,28 @@
 HighLevelType: TypeAlias = "Callable[..., T]"
 
 
-def high_level_function(func: DispatcherType) -> HighLevelType:
+def high_level_function(
+    module: str = "ak", name: str | None = None
+) -> Callable[[DispatcherType], HighLevelType]:
     """Decorate a high-level function such that it may be overloaded by third-party array objects"""
 
+    def capture_func(func: DispatcherType) -> HighLevelType:
+        if name is None:
+            captured_name = func.__qualname__
+        else:
+            captured_name = name
+        return named_high_level_function(func, f"{module}.{captured_name}")
+
+    return capture_func
+
+
+def named_high_level_function(func: DispatcherType, name: str) -> HighLevelType:
+    """Decorate a named high-level function such that it may be overloaded by third-party array objects"""
+
     @wraps(func)
     def dispatch(*args, **kwargs):
         # NOTE: this decorator assumes that the operation is exposed under `ak.`
-        with OperationErrorContext(f"ak.{func.__qualname__}", args, kwargs):
+        with OperationErrorContext(name, args, kwargs):
             gen_or_result = func(*args, **kwargs)
             if isgenerator(gen_or_result):
                 array_likes = next(gen_or_result)

src/awkward/operations/ak_all.py

@@ -11,7 +11,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def all(
     array,
     axis=None,

src/awkward/operations/ak_almost_equal.py

@@ -12,7 +12,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def almost_equal(
     left,
     right,

src/awkward/operations/ak_any.py

@@ -11,7 +11,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def any(
     array,
     axis=None,

src/awkward/operations/ak_argcartesian.py

@@ -15,7 +15,7 @@
 cpu = NumpyBackend.instance()
 
 
-@high_level_function
+@high_level_function()
 def argcartesian(
     arrays,
     axis=1,

src/awkward/operations/ak_argcombinations.py

@@ -9,7 +9,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def argcombinations(
     array,
     n,

src/awkward/operations/ak_argmax.py

@@ -11,7 +11,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def argmax(
     array,
     axis=None,
@@ -67,7 +67,7 @@ def argmax(
     return _impl(array, axis, keepdims, mask_identity, highlevel, behavior)
 
 
-@high_level_function
+@high_level_function()
 def nanargmax(
     array,
     axis=None,

src/awkward/operations/ak_argmin.py

@@ -11,7 +11,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def argmin(
     array,
     axis=None,
@@ -67,7 +67,7 @@ def argmin(
     return _impl(array, axis, keepdims, mask_identity, highlevel, behavior)
 
 
-@high_level_function
+@high_level_function()
 def nanargmin(
     array,
     axis=None,

src/awkward/operations/ak_argsort.py

@@ -10,7 +10,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def argsort(
     array, axis=-1, *, ascending=True, stable=True, highlevel=True, behavior=None
 ):

src/awkward/operations/ak_backend.py

@@ -4,7 +4,7 @@
 from awkward._dispatch import high_level_function
 
 
-@high_level_function
+@high_level_function()
 def backend(*arrays):
     """
     Args:

src/awkward/operations/ak_broadcast_arrays.py

@@ -13,7 +13,7 @@
 cpu = NumpyBackend.instance()
 
 
-@high_level_function
+@high_level_function()
 def broadcast_arrays(
     *arrays,
     depth_limit=None,

src/awkward/operations/ak_broadcast_fields.py

@@ -10,7 +10,7 @@
 cpu = NumpyBackend.instance()
 
 
-@high_level_function
+@high_level_function()
 def broadcast_fields(*arrays, highlevel=True, behavior=None):
     """
     Args:

src/awkward/operations/ak_cartesian.py

@@ -16,7 +16,7 @@
 cpu = NumpyBackend.instance()
 
 
-@high_level_function
+@high_level_function()
 def cartesian(
     arrays,
     axis=1,

src/awkward/operations/ak_categories.py

@@ -6,7 +6,7 @@
 from awkward._layout import wrap_layout
 
 
-@high_level_function
+@high_level_function()
 def categories(array, highlevel=True):
     """
     Args:

src/awkward/operations/ak_combinations.py

@@ -9,7 +9,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def combinations(
     array,
     n,

src/awkward/operations/ak_concatenate.py

@@ -18,7 +18,7 @@
 
 
 @ak._connect.numpy.implements("concatenate")
-@high_level_function
+@high_level_function()
 def concatenate(arrays, axis=0, *, mergebool=True, highlevel=True, behavior=None):
     """
     Args:

src/awkward/operations/ak_copy.py

@@ -10,7 +10,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def copy(array):
     """
     Args:

src/awkward/operations/ak_corr.py

@@ -10,7 +10,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def corr(x, y, weight=None, axis=None, *, keepdims=False, mask_identity=False):
     """
     Args:

src/awkward/operations/ak_count.py

@@ -10,7 +10,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def count(
     array,
     axis=None,

src/awkward/operations/ak_count_nonzero.py

@@ -10,7 +10,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def count_nonzero(
     array,
     axis=None,

src/awkward/operations/ak_covar.py

@@ -9,7 +9,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def covar(x, y, weight=None, axis=None, *, keepdims=False, mask_identity=False):
     """
     Args:

src/awkward/operations/ak_drop_none.py

@@ -10,7 +10,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def drop_none(array, axis=None, highlevel=True, behavior=None):
     """
     Args:

src/awkward/operations/ak_enforce_type.py

@@ -17,7 +17,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def enforce_type(array, type, *, highlevel=True, behavior=None):
     """
     Args:

src/awkward/operations/ak_fields.py

@@ -7,7 +7,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def fields(array):
     """
     Args:

src/awkward/operations/ak_fill_none.py

@@ -16,7 +16,7 @@
 cpu = NumpyBackend.instance()
 
 
-@high_level_function
+@high_level_function()
 def fill_none(array, value, axis=-1, *, highlevel=True, behavior=None):
     """
     Args:

src/awkward/operations/ak_firsts.py

@@ -11,7 +11,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def firsts(array, axis=1, *, highlevel=True, behavior=None):
     """
     Args:

src/awkward/operations/ak_flatten.py

@@ -9,7 +9,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def flatten(array, axis=1, *, highlevel=True, behavior=None):
     """
     Args:

src/awkward/operations/ak_from_arrow.py

@@ -8,7 +8,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def from_arrow(array, *, generate_bitmasks=False, highlevel=True, behavior=None):
     """
     Args:

src/awkward/operations/ak_from_arrow_schema.py

@@ -6,7 +6,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def from_arrow_schema(schema):
     """
     Args:

src/awkward/operations/ak_from_avro_file.py

@@ -10,7 +10,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def from_avro_file(
     file, limit_entries=None, *, debug_forth=False, highlevel=True, behavior=None
 ):

src/awkward/operations/ak_from_buffers.py

@@ -16,7 +16,7 @@
 numpy = Numpy.instance()
 
 
-@high_level_function
+@high_level_function()
 def from_buffers(
     form,
     length,

src/awkward/operations/ak_from_categorical.py

@@ -6,7 +6,7 @@
 from awkward._layout import wrap_layout
 
 
-@high_level_function
+@high_level_function()
 def from_categorical(array, *, highlevel=True, behavior=None):
     """
     Args:

src/awkward/operations/ak_from_cupy.py

@@ -4,7 +4,7 @@
 from awkward._layout import from_arraylib, wrap_layout
 
 
-@high_level_function
+@high_level_function()
 def from_cupy(array, *, regulararray=False, highlevel=True, behavior=None):
     """
     Args:

src/awkward/operations/ak_from_iter.py

@@ -11,7 +11,7 @@
 np = NumpyMetadata.instance()
 
 
-@high_level_function
+@high_level_function()
 def from_iter(
     iterable,
     *,

src/awkward/operations/ak_from_jax.py

@@ -5,7 +5,7 @@
 from awkward._layout import from_arraylib, wrap_layout
 
 
-@high_level_function
+@high_level_function()
 def from_jax(array, *, regulararray=False, highlevel=True, behavior=None):
     """
     Args:

src/awkward/operations/ak_from_json.py

@@ -18,7 +18,7 @@
 numpy = Numpy.instance()
 
 
-@high_level_function
+@high_level_function()
 def from_json(
     source,
     *,

src/awkward/operations/ak_from_numpy.py

@@ -4,7 +4,7 @@
 from awkward._layout import from_arraylib, wrap_layout
 
 
-@high_level_function
+@high_level_function()
 def from_numpy(
     array, *, regulararray=False, recordarray=True, highlevel=True, behavior=None
 ):

src/awkward/operations/ak_from_parquet.py

@@ -6,7 +6,7 @@
 from awkward._regularize import is_integer
 
 
-@high_level_function
+@high_level_function()
 def from_parquet(
     path,
     *,
@@ -76,7 +76,7 @@ def from_parquet(
     )
 
 
-@high_level_function
+@high_level_function()
 def metadata(
     path,
     storage_options=None,