feature: add type hints for recipes (pytoolz#496)

.gitignore

@@ -4,5 +4,6 @@ dist/
 *.egg-info/
 bench/shakespeare.txt
 .coverage
+.idea
 
 \.tox/

.travis.yml

@@ -18,6 +18,7 @@ env:
 # command to install dependencies
 install:
     - pip install coverage pep8 pytest
+    - if [[ $TRAVIS_PYTHON_VERSION != pypy* ]]; then pip install 'mypy==0.790' ; fi
 
 # command to run tests
 # require 100% coverage (not including test files) to pass Travis CI test
@@ -32,6 +33,7 @@ script:
     - python setup.py develop
     - py.test
     - nosetests
+    - if [[ $TRAVIS_PYTHON_VERSION != pypy* ]]; then mypy toolz/recipes.py ; fi
 
 # load coverage status to https://coveralls.io
 after_success:

toolz/__init__.py

@@ -19,7 +19,7 @@
 
 from . import curried, sandbox
 
-functoolz._sigs.create_signature_registry()
+functoolz._sigs.create_signature_registry()  # type: ignore
 
 from ._version import get_versions
 __version__ = get_versions()['version']

toolz/_signatures.py

@@ -571,7 +571,7 @@
         lambda a, b: None],
 )
 
-module_info['toolz'] = dict(
+module_info['toolz'] = dict(  # type: ignore
     curry=[
         (0, lambda *args, **kwargs: None)],
     excepts=[
@@ -584,7 +584,7 @@
         (0, lambda func=None, cache=None, key=None: None)],
 )
 
-module_info['toolz.functoolz'] = dict(
+module_info['toolz.functoolz'] = dict(  # type: ignore
     Compose=[
         (0, lambda funcs: None)],
     InstanceProperty=[
@@ -653,7 +653,7 @@ def expand_sig(sig):
     return num_pos_only, func, keyword_only + keyword_exclude, sigspec
 
 
-signatures = {}
+signatures = {}  # type: ignore
 
 
 def create_signature_registry(module_info=module_info, signatures=signatures):

toolz/_version.py

@@ -52,7 +52,7 @@ class NotThisMethod(Exception):
     """Exception raised if a method is not valid for the current scenario."""
 
 
-LONG_VERSION_PY = {}
+LONG_VERSION_PY = {}  # type: ignore
 HANDLERS = {}
 
 

toolz/curried/__init__.py

@@ -99,5 +99,5 @@
 valfilter = toolz.curry(toolz.valfilter)
 valmap = toolz.curry(toolz.valmap)
 
-del exceptions
+del exceptions  # type: ignore
 del toolz

toolz/recipes.py

@@ -1,11 +1,36 @@
 import itertools
-from .itertoolz import frequencies, pluck, getter
+from typing import TypeVar, Dict, Tuple, Callable, Any, Iterable, Sequence, \
+    overload, List
 
+from .itertoolz import frequencies, pluck, getter
 
 __all__ = ('countby', 'partitionby')
 
+A = TypeVar('A')
+B = TypeVar('B')
+KeyLike = TypeVar('KeyLike', int, Iterable, Callable, Tuple)
+
+
+# Case: countby(len, ['cat', 'mouse', 'dog'])
+@overload
+def countby(key: Callable[[A], B], seq: Iterable[A]) -> Dict[B, int]: ...
+
+
+# Case: countby('a', [{'a': 1, 'b': 2}, {'a': 10, 'b': 2}])
+@overload
+def countby(key: A, seq: Iterable[Dict[A, B]]) -> Dict[B, int]: ...
+
+
+# Case: countby(0, [[1, 2], [10, 2]])
+def countby(key: int, seq: Iterable[Iterable[A]]) -> Dict[A, int]: ...
+
+
+# Case: countby([0, 1], [[1, 2], [10, 2]])
+def countby(key: List[int],
+            seq: Iterable[Iterable[A]]) -> Dict[Tuple[A, ...], int]: ...
+
 
-def countby(key, seq):
+def countby(key: KeyLike, seq: Iterable[Any]) -> Dict[Any, int]:
     """ Count elements of a collection by a key function
 
     >>> countby(len, ['cat', 'mouse', 'dog'])
@@ -20,10 +45,11 @@ def countby(key, seq):
     """
     if not callable(key):
         key = getter(key)
-    return frequencies(map(key, seq))
+    return frequencies(map(key, seq))  # type: ignore
 
 
-def partitionby(func, seq):
+def partitionby(func: Callable[[A], bool],
+                seq: Sequence[A]) -> Iterable[Tuple[A, ...]]:
     """ Partition a sequence according to a function
 
     Partition `s` into a sequence of lists such that, when traversing