From 2d91d8532ab7ab5698f4e030ee558beb3cd124f5 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 +
 toolz/recipes.py | 11 ++++++++---
 2 files changed, 9 insertions(+), 3 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/toolz/recipes.py b/toolz/recipes.py
index 89de88db..fdd5eb26 100644
--- a/toolz/recipes.py
+++ b/toolz/recipes.py
@@ -1,11 +1,15 @@
 import itertools
-from .itertoolz import frequencies, pluck, getter
+from typing import TypeVar, Iterable, Dict, Tuple, Callable
 
+from .itertoolz import frequencies, pluck, getter
 
 __all__ = ('countby', 'partitionby')
 
+A = TypeVar('A')
+B = TypeVar('B')
+
 
-def countby(key, seq):
+def countby(key: Callable[[A], B], seq: Iterable[A]) -> Dict[B, int]:
     """ Count elements of a collection by a key function
 
     >>> countby(len, ['cat', 'mouse', 'dog'])
@@ -23,7 +27,8 @@ def countby(key, seq):
     return frequencies(map(key, seq))
 
 
-def partitionby(func, seq):
+def partitionby(func: Callable[[A], bool],
+                seq: Iterable[A]) -> Iterable[Tuple[A, ...]]:
     """ Partition a sequence according to a function
 
     Partition `s` into a sequence of lists such that, when traversing