From 1840a2de473e36c4c23f63d0734e6846e884ebb5 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               |  4 +++-
 toolz/__init__.py         |  2 +-
 toolz/_signatures.py      |  6 +++---
 toolz/_version.py         |  2 +-
 toolz/curried/__init__.py |  2 +-
 toolz/recipes.py          | 14 ++++++++++----
 7 files changed, 20 insertions(+), 11 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..3a2e0ae8 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -17,7 +17,8 @@ env:
 
 # command to install dependencies
 install:
-    - pip install coverage pep8 pytest
+    - sudo apt-get install python3-dev pypy-dev # Needed for pypy compilation of mypy
+    - pip install coverage pep8 pytest 'mypy==0.782'
 
 # 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
+    - mypy toolz/recipes.py
 
 # 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..29f5bf6b 100644
--- a/toolz/recipes.py
+++ b/toolz/recipes.py
@@ -1,11 +1,16 @@
 import itertools
-from .itertoolz import frequencies, pluck, getter
+from typing import TypeVar, Dict, Tuple, Callable, Any, Iterable, Sequence
 
+from .itertoolz import frequencies, pluck, getter
 
 __all__ = ('countby', 'partitionby')
 
+A = TypeVar('A')
+B = TypeVar('B')
+KeyLike = TypeVar('KeyLike', int, Iterable, Callable, Tuple)
+
 
-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 +25,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