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