From 646ef40015c1cadfe2aca2ca807c822fd4da17b7 Mon Sep 17 00:00:00 2001
From: Jiehong Ma <email@majiehong.com>
Date: Thu, 8 Oct 2020 10:14:57 +0200
Subject: [PATCH] feature: add type hints for recipes (#496)

---
 .gitignore                |  1 +
 .travis.yml               |  2 ++
 toolz/__init__.py         |  2 +-
 toolz/_signatures.py      |  6 +++---
 toolz/_version.py         |  2 +-
 toolz/curried/__init__.py |  2 +-
 toolz/recipes.py          | 34 ++++++++++++++++++++++++++++++----
 7 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/.gitignore b/.gitignore
index cfef783b..afd39a7f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,5 +4,6 @@ dist/
 *.egg-info/
 bench/shakespeare.txt
 .coverage
+.idea
 
 \.tox/
diff --git a/.travis.yml b/.travis.yml
index 8d496be0..74938585 100644
--- a/.travis.yml
+++ b/.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:
diff --git a/toolz/__init__.py b/toolz/__init__.py
index ba49a662..273f4ac0 100644
--- a/toolz/__init__.py
+++ b/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']
diff --git a/toolz/_signatures.py b/toolz/_signatures.py
index 328cea91..603b1d32 100644
--- a/toolz/_signatures.py
+++ b/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):
diff --git a/toolz/_version.py b/toolz/_version.py
index 7d6792fd..30cf0e74 100644
--- a/toolz/_version.py
+++ b/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 = {}
 
 
diff --git a/toolz/curried/__init__.py b/toolz/curried/__init__.py
index 356eddbd..2a46dd16 100644
--- a/toolz/curried/__init__.py
+++ b/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
diff --git a/toolz/recipes.py b/toolz/recipes.py
index 89de88db..3e5a91de 100644
--- a/toolz/recipes.py
+++ b/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