TYP: annotate curry

toolz/functoolz.py

@@ -1,22 +1,26 @@
-from functools import reduce, partial
+from __future__ import annotations
+
 import inspect
 import sys
-from operator import attrgetter, not_
+from functools import partial, reduce
 from importlib import import_module
+from operator import attrgetter, not_
 from textwrap import dedent
 from types import MethodType
+from typing import Any, Callable, Dict, Generic, TypeVar, Union, overload
 
 from .utils import no_default
 
-
 __all__ = ('identity', 'apply', 'thread_first', 'thread_last', 'memoize',
            'compose', 'compose_left', 'pipe', 'complement', 'juxt', 'do',
            'curry', 'flip', 'excepts')
 
 PYPY = hasattr(sys, 'pypy_version_info')
+_T = TypeVar("_T")
+_T2 = TypeVar("_T2")
 
 
-def identity(x):
+def identity(x: _T) ->_T:
     """ Identity function. Return x
 
     >>> identity(3)
@@ -164,7 +168,7 @@ def __reduce__(self):
         return InstanceProperty, state
 
 
-class curry(object):
+class curry(Generic[_T]):
     """ Curry a callable function
 
     Enables partial application of arguments through calling a function with an
@@ -192,7 +196,14 @@ class curry(object):
         toolz.curried - namespace of curried functions
                         https://toolz.readthedocs.io/en/latest/curry.html
     """
-    def __init__(self, *args, **kwargs):
+    @overload
+    def __init__(self, func: Callable[..., _T], *args: Any, **kwargs: Any) -> None: ...
+
+    # this overload should never be used, mypy complains if only one overload exists
+    @overload
+    def __init__(self, *args: Union[Callable[..., _T], Any], **kwargs: Any) -> None: ...
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         if not args:
             raise TypeError('__init__() takes at least 2 arguments (1 given)')
         func, args = args[0], args[1:]
@@ -227,7 +238,7 @@ def __init__(self, *args, **kwargs):
         self._has_unknown_args = None
 
     @instanceproperty
-    def func(self):
+    def func(self) -> Callable[..., _T]:
         return self._partial.func
 
     @instanceproperty
@@ -273,32 +284,32 @@ def args(self):
         return self._partial.args
 
     @instanceproperty
-    def keywords(self):
+    def keywords(self) -> Dict[str, Any]:
         return self._partial.keywords
 
     @instanceproperty
-    def func_name(self):
+    def func_name(self) -> str:
         return self.__name__
 
-    def __str__(self):
+    def __str__(self) -> str:
         return str(self.func)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return repr(self.func)
 
-    def __hash__(self):
+    def __hash__(self) -> int:
         return hash((self.func, self.args,
                      frozenset(self.keywords.items()) if self.keywords
                      else None))
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         return (isinstance(other, curry) and self.func == other.func and
                 self.args == other.args and self.keywords == other.keywords)
 
-    def __ne__(self, other):
+    def __ne__(self, other: Any) -> bool:
         return not self.__eq__(other)
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: Any, **kwargs: Any) -> Union[_T, curry[_T]]:
         try:
             return self._partial(*args, **kwargs)
         except TypeError as exc:
@@ -332,10 +343,10 @@ def _should_curry(self, args, kwargs, exc=None):
             # There was a genuine TypeError
             return False
 
-    def bind(self, *args, **kwargs):
+    def bind(self, *args, **kwargs) -> curry[_T]:
         return type(self)(self, *args, **kwargs)
 
-    def call(self, *args, **kwargs):
+    def call(self, *args, **kwargs) -> _T:
         return self._partial(*args, **kwargs)
 
     def __get__(self, instance, owner):
@@ -677,7 +688,7 @@ def __setstate__(self, state):
         self.funcs = state
 
 
-def do(func, x):
+def do(func: Callable[[_T], Any], x: _T) -> _T:
     """ Runs ``func`` on ``x``, returns ``x``
 
     Because the results of ``func`` are not returned, only the side
@@ -704,7 +715,7 @@ def do(func, x):
 
 
 @curry
-def flip(func, a, b):
+def flip(func: Callable[[_T, _T], _T2], a: _T, b: _T) -> _T2:
     """ Call the function call with the arguments flipped
 
     This function is curried.
@@ -730,7 +741,7 @@ def flip(func, a, b):
     return func(b, a)
 
 
-def return_none(exc):
+def return_none(exc: Any) -> None:
     """ Returns None.
     """
     return None