typing: improve applicator overloads

eth_utils/applicators.py

@@ -1,9 +1,13 @@
 from collections.abc import (
     Callable,
     Generator,
+    Sequence,
 )
 from typing import (
     Any,
+    TypeGuard,
+    TypeVar,
+    overload,
 )
 import warnings
 
@@ -22,12 +26,15 @@
 )
 
 Formatters = Callable[[list[Any]], list[Any]]
+TArgument = TypeVar("TArgument")
+TOther = TypeVar("TOther")
+TReturn = TypeVar("TReturn")
 
 
 @return_arg_type(2)
 def apply_formatter_at_index(
-    formatter: Callable[..., Any], at_index: int, value: list[Any]
-) -> Generator[list[Any], None, None]:
+    formatter: Callable[[Any], TReturn], at_index: int, value: Sequence[Any]
+) -> Generator[Any | TReturn, None, None]:
     if at_index + 1 > len(value):
         raise IndexError(
             f"Not enough values in iterable to apply formatter. Got: {len(value)}. "
@@ -40,7 +47,7 @@ def apply_formatter_at_index(
             yield item
 
 
-def combine_argument_formatters(*formatters: list[Callable[..., Any]]) -> Formatters:
+def combine_argument_formatters(*formatters: Callable[[Any], Any]) -> Formatters:
     warnings.warn(
         DeprecationWarning(
             "combine_argument_formatters(formatter1, formatter2)([item1, item2])"
@@ -63,8 +70,8 @@ def combine_argument_formatters(*formatters: list[Callable[..., Any]]) -> Format
 
 @return_arg_type(1)
 def apply_formatters_to_sequence(
-    formatters: list[Any], sequence: list[Any]
-) -> Generator[list[Any], None, None]:
+    formatters: Sequence[Callable[[Any], TReturn]], sequence: Sequence[Any]
+) -> Generator[TReturn, None, None]:
     if len(formatters) == len(sequence):
         for formatter, item in zip(formatters, sequence):
             yield formatter(item)
@@ -80,9 +87,29 @@ def apply_formatters_to_sequence(
         )
 
 
+@overload
 def apply_formatter_if(
-    condition: Callable[..., bool], formatter: Callable[..., Any], value: Any
-) -> Any:
+    condition: Callable[[TArgument], TypeGuard[TOther]],
+    formatter: Callable[[TOther], TReturn],
+    value: TArgument,
+) -> TArgument | TReturn:
+    ...
+
+
+@overload
+def apply_formatter_if(
+    condition: Callable[[TArgument], bool],
+    formatter: Callable[[TArgument], TReturn],
+    value: TArgument,
+) -> TArgument | TReturn:
+    ...
+
+
+def apply_formatter_if(
+    condition: Callable[[TArgument], bool],
+    formatter: Callable[[Any], TReturn],
+    value: TArgument,
+) -> TArgument | TReturn:
     if condition(value):
         return formatter(value)
     else:
@@ -129,16 +156,18 @@ def apply_formatters_to_dict(
 
 @return_arg_type(1)
 def apply_formatter_to_array(
-    formatter: Callable[..., Any], value: list[Any]
-) -> Generator[list[Any], None, None]:
+    formatter: Callable[[TArgument], TReturn], value: Sequence[TArgument]
+) -> Generator[TReturn, None, None]:
     for item in value:
         yield formatter(item)
 
 
 def apply_one_of_formatters(
-    formatter_condition_pairs: tuple[tuple[Callable[..., Any], Callable[..., Any]]],
-    value: Any,
-) -> Any:
+    formatter_condition_pairs: Sequence[
+        tuple[Callable[[TArgument], bool], Callable[[TArgument], TReturn]]
+    ],
+    value: TArgument,
+) -> TReturn:
     for condition, formatter in formatter_condition_pairs:
         if condition(value):
             return formatter(value)

eth_utils/curried/__init__.py

@@ -6,6 +6,7 @@
 from typing import (
     Any,
     Optional,
+    TypeGuard,
     TypeVar,
     Union,
     overload,
@@ -119,69 +120,100 @@
 )
 
 TReturn = TypeVar("TReturn")
-TValue = TypeVar("TValue")
+TArgument = TypeVar("TArgument")
+TOther = TypeVar("TOther")
 
 
 @overload
 def apply_formatter_if(
-    condition: Callable[..., bool],
-) -> Callable[[Callable[..., TReturn]], Callable[[TValue], TReturn | TValue]]:
+    condition: Callable[[TArgument], TypeGuard[TOther]],
+) -> Callable[
+    [Callable[[TOther], TReturn]], Callable[[TArgument], TArgument | TReturn]
+]:
+    pass
+
+
+@overload
+def apply_formatter_if(
+    condition: Callable[[TArgument], bool],
+) -> Callable[
+    [Callable[[TArgument], TReturn]], Callable[[TArgument], TArgument | TReturn]
+]:
+    pass
+
+
+@overload
+def apply_formatter_if(
+    condition: Callable[[TArgument], TypeGuard[TOther]],
+    formatter: Callable[[TOther], TReturn],
+) -> Callable[[TArgument], TArgument | TReturn]:
+    pass
+
+
+@overload
+def apply_formatter_if(
+    condition: Callable[[TArgument], bool], formatter: Callable[[TArgument], TReturn]
+) -> Callable[[TArgument], TArgument | TReturn]:
     pass
 
 
 @overload
 def apply_formatter_if(
-    condition: Callable[..., bool], formatter: Callable[..., TReturn]
-) -> Callable[[TValue], TReturn | TValue]:
+    condition: Callable[[TArgument], TypeGuard[TOther]],
+    formatter: Callable[[TOther], TReturn],
+    value: TArgument,
+) -> TArgument | TReturn:
     pass
 
 
 @overload
 def apply_formatter_if(
-    condition: Callable[..., bool], formatter: Callable[..., TReturn], value: TValue
-) -> TReturn | TValue:
+    condition: Callable[[TArgument], bool],
+    formatter: Callable[[TArgument], TReturn],
+    value: TArgument,
+) -> TArgument | TReturn:
     pass
 
 
 # This is just a stub to appease mypy, it gets overwritten later
 def apply_formatter_if(  # type: ignore
     condition: Callable[..., bool],
     formatter: Callable[..., TReturn] | None = None,
-    value: TValue | None = None,
+    value: TArgument | None = None,
 ) -> (
-    Callable[[Callable[..., TReturn]], Callable[[TValue], TReturn | TValue]]
-    | Callable[[TValue], TReturn | TValue]
+    Callable[[Callable[..., TReturn]], Callable[[TArgument], TReturn | TArgument]]
+    | Callable[[TArgument], TReturn | TArgument]
     | TReturn
-    | TValue
+    | TArgument
 ):
     pass
 
 
 @overload
 def apply_one_of_formatters(
     formatter_condition_pairs: Sequence[
-        tuple[Callable[..., bool], Callable[..., TReturn]]
+        tuple[Callable[[TArgument], bool], Callable[[TArgument], TReturn]]
     ],
-) -> Callable[[TValue], TReturn]:
+) -> Callable[[TArgument], TReturn]:
     ...
 
 
 @overload
 def apply_one_of_formatters(
     formatter_condition_pairs: Sequence[
-        tuple[Callable[..., bool], Callable[..., TReturn]]
+        tuple[Callable[[TArgument], bool], Callable[[TArgument], TReturn]]
     ],
-    value: TValue,
+    value: TArgument,
 ) -> TReturn:
     ...
 
 
 # This is just a stub to appease mypy, it gets overwritten later
 def apply_one_of_formatters(  # type: ignore
     formatter_condition_pairs: Sequence[
-        tuple[Callable[..., bool], Callable[..., TReturn]]
+        tuple[Callable[[TArgument], bool], Callable[[TArgument], TReturn]]
     ],
-    value: TValue | None = None,
+    value: TArgument | None = None,
 ) -> TReturn:
     ...
 
@@ -282,8 +314,10 @@ def apply_formatters_to_dict(  # type: ignore
 del Generator
 del Optional
 del Sequence
+del TArgument
+del TOther
 del TReturn
-del TValue
+del TypeGuard
 del TypeVar
 del Union
 del curry

tests/mypy/applicator_overloads.py

@@ -0,0 +1,76 @@
+# no-unused-vars
+from collections.abc import (
+    Callable,
+)
+from typing import (
+    Any,
+    TypeGuard,
+)
+
+from eth_utils import (
+    apply_formatter_at_index,
+    apply_formatter_if,
+    apply_formatter_to_array,
+    apply_formatters_to_sequence,
+    apply_one_of_formatters,
+    combine_argument_formatters,
+)
+from eth_utils.curried import (
+    apply_formatter_if as curried_apply_formatter_if,
+    apply_one_of_formatters as curried_apply_one_of_formatters,
+)
+
+
+def is_text_value(value: str | int) -> TypeGuard[str]:
+    return isinstance(value, str)
+
+
+def is_text_condition(value: str) -> bool:
+    return isinstance(value, str)
+
+
+def uppercase(value: str) -> str:
+    return value.upper()
+
+
+def test_sequence_formatter_inputs() -> None:
+    array_result: tuple[int, ...] = apply_formatter_to_array(  # noqa: F841
+        int, ("1", "2")
+    )
+    sequence_result: tuple[int, ...] = apply_formatters_to_sequence(  # noqa: F841
+        (int, int), ("1", "2")
+    )
+    index_result: tuple[int | str, ...] = apply_formatter_at_index(  # noqa: F841
+        str, 1, (1, 2, 3)
+    )
+    combined: Callable[  # noqa: F841
+        [list[Any]], list[Any]
+    ] = combine_argument_formatters(str, int)
+
+
+def test_condition_formatter_inputs() -> None:
+    value: str | int = "abc"
+    guarded_result: str | int = apply_formatter_if(  # noqa: F841
+        is_text_value, uppercase, value
+    )
+    bool_result: str | int = apply_formatter_if(  # noqa: F841
+        is_text_condition, len, "abc"
+    )
+    curried_result: Callable[  # noqa: F841
+        [str | int], str | int
+    ] = curried_apply_formatter_if(is_text_value, uppercase)
+    condition_first = curried_apply_formatter_if(is_text_value)
+    formatter_first: Callable[[str | int], str | int] = condition_first(  # noqa: F841
+        uppercase
+    )
+
+
+def test_one_of_formatter_inputs() -> None:
+    formatter_pairs: tuple[tuple[Callable[[str], bool], Callable[[str], str]], ...] = (
+        (is_text_condition, uppercase),
+    )
+    direct_result: str = apply_one_of_formatters(formatter_pairs, "abc")  # noqa: F841
+    curried_result: Callable[  # noqa: F841
+        [str], str
+    ] = curried_apply_one_of_formatters(formatter_pairs)
+    curried_value: str = curried_result("abc")  # noqa: F841