feat[next]: extend embedded implementation of premap()

src/gt4py/_core/definitions.py

@@ -20,7 +20,6 @@
 import functools
 import math
 import numbers
-from typing import overload
 
 import numpy as np
 import numpy.typing as npt
@@ -42,6 +41,7 @@
     TypeVar,
     Union,
     cast,
+    overload,
 )
 
 

src/gt4py/eve/extended_typing.py

@@ -127,11 +127,16 @@ def __dir__() -> List[str]:
     return self_func.__cached_dir
 
 
-_T = TypeVar("_T")
-
 # -- Common type aliases --
 NoArgsCallable = Callable[[], Any]
 
+_A = TypeVar("_A", contravariant=True)
+_R = TypeVar("_R", covariant=True)
+
+
+class ArgsOnlyCallable(Protocol[_A, _R]):
+    def __call__(self, *args: _A) -> _R: ...
+
 
 # -- Typing annotations --
 if _sys.version_info >= (3, 9):
@@ -367,6 +372,9 @@ def has_type_parameters(cls: Type) -> bool:
     return issubclass(cls, Generic) and len(getattr(cls, "__parameters__", [])) > 0  # type: ignore[arg-type]  # Generic not considered as a class
 
 
+_T = TypeVar("_T")
+
+
 def get_actual_type(obj: _T) -> Type[_T]:
     """Return type of an object (also working for GenericAlias instances which pretend to be an actual type)."""
     return StdGenericAliasType if isinstance(obj, StdGenericAliasType) else type(obj)

src/gt4py/eve/utils.py

@@ -52,6 +52,7 @@
 from . import extended_typing as xtyping
 from .extended_typing import (
     Any,
+    ArgsOnlyCallable,
     Callable,
     Collection,
     Dict,
@@ -84,6 +85,15 @@
 T = TypeVar("T")
 
 
+def first(iterable: Iterable[T], *, default: Union[T, NothingType] = NOTHING) -> T:
+    try:
+        return next(iter(iterable))
+    except StopIteration as error:
+        if default is not NOTHING:
+            return cast(T, default)
+        raise error
+
+
 def isinstancechecker(type_info: Union[Type, Iterable[Type]]) -> Callable[[Any], bool]:
     """Return a callable object that checks if operand is an instance of `type_info`.
 
@@ -227,9 +237,31 @@ def itemgetter_(key: Any, default: Any = NOTHING) -> Callable[[Any], Any]:
 
 
 _P = ParamSpec("_P")
+_S = TypeVar("_S")
 _T = TypeVar("_T")
 
 
+@dataclasses.dataclass(frozen=True)
+class IndexerCallable(Generic[_S, _T]):
+    """
+    An indexer class applying the wrapped function to the index arguments.
+
+    Examples:
+        >>> indexer = IndexerCallable(lambda x: x**2)
+        >>> indexer[3]
+        9
+
+        >>> indexer = IndexerCallable(lambda a, b: a + b)
+        >>> indexer[3, 4]
+        7
+    """
+
+    func: ArgsOnlyCallable[_S, _T]
+
+    def __getitem__(self, key: _S | Tuple[_S, ...]) -> _T:
+        return self.func(*key) if isinstance(key, tuple) else self.func(key)
+
+
 class fluid_partial(functools.partial):
     """Create a `functools.partial` with support for multiple applications calling `.partial()`."""