added mypy type annotation checking

.travis.yml

@@ -12,6 +12,7 @@ install:
   - pip install -r requirements.txt
   - pip install coveralls
 script:
+  - mypy basic_utils/*.py
   - coverage run --branch -m pytest && coverage report -m
   - python -m doctest basic_utils/[a-zA-Z]*.py
 after_success:

basic_utils/core.py

@@ -1,13 +1,19 @@
 from collections import defaultdict
 from functools import reduce
+from typing import List
 
 import operator
 import os
 
 sentinel = object()
 
+__all__ = [
+    'clear', 'getattrs', 'map_getattr', 'recursive_default_dict', 'rgetattr',
+    'rsetattr', 'slurp', 'to_string',
+]
 
-def slurp(fname):
+
+def slurp(fname: str) -> str:
     """
     Reads a file and all its contents, returns a single string
     """
@@ -16,14 +22,14 @@ def slurp(fname):
     return data
 
 
-def clear():
+def clear() -> None:
     """
     Clears the terminal screen from python, operating system agnostic
     """
     os.system('cls' if os.name == 'nt' else 'clear')
 
 
-def to_string(objects, sep=", "):
+def to_string(objects: List[object], sep: str = ", ") -> str:
     """
     Converts a list of objects into a single string
 
@@ -59,14 +65,3 @@ def rsetattr(obj, attr, val):
     """Sets a nested attribute within an object"""
     pre, _, post = attr.rpartition('.')
     return setattr(rgetattr(obj, pre) if pre else obj, post, val)
-
-
-def identity(*args):
-    """
-    Returns the same values passed as arguments
-
-    >>> identity(10, 20)
-    (10, 20)
-    """
-    first, *rest = args
-    return first if not rest else args

basic_utils/helpers.py

@@ -1,10 +1,10 @@
-from collections import deque, abc, Iterable
 from functools import reduce
 from itertools import chain
 from operator import getitem
+from typing import Sequence, Iterable, Any, Callable
 
 
-def first(seq):
+def first(seq: Sequence) -> Any:
     """
     Returns first element in a sequence.
 
@@ -14,70 +14,57 @@ def first(seq):
     return next(iter(seq))
 
 
-def rest(seq):
+def rest(seq: Sequence) -> Any:
     """
     Returns remaining elements in a sequence
 
     >>> rest([1, 2, 3])
     [2, 3]
     """
-    if isinstance(seq, (abc.Sequence, str)):
-        return seq[1:]
-    d = deque(seq)
-    d.popleft()
-    return tuple(d)
+    return seq[1:]
 
 
-def last(seq):
+def last(seq: Sequence) -> Any:
     """
     Returns the last item in a sequence
 
     >>> last([1, 2, 3])
     3
     """
-    if hasattr(seq, '__getitem__'):
-        return seq[-1]
-    return deque(iter(seq), maxlen=1).pop()
+    return seq[-1]
 
 
-def butlast(seq):
+def butlast(seq: Sequence) -> Any:
     """
     Returns an iterable of all but the last item in the sequence
 
     >>> butlast([1, 2, 3])
     [1, 2]
     """
-    if isinstance(seq, (abc.Sequence, str)):
-        return seq[:-1]
-    i = deque(seq)
-    i.pop()
-    return tuple(i)
+    return seq[:-1]
 
 
-def reverse(seq):
+def reverse(seq: Sequence) -> Sequence:
     """
     Returns a sequence of items in seq in reverse order
 
     >>> reverse([1, 2, 3])
     [3, 2, 1]
     """
-    try:
-        return seq[::-1]
-    except TypeError:
-        return tuple(y for y in reverse(list(seq)))
+    return seq[::-1]
 
 
-def partial_flatten(seq):
+def partial_flatten(seq: Iterable) -> Iterable:
     """
     Returns partially flattened version of seq
 
     >>> list(flatten([[1, 2, 3], [4, 5, 6]]))
     [1, 2, 3, 4, 5, 6]
     """
-    return type(seq)(chain.from_iterable(seq))
+    return chain.from_iterable(seq)
 
 
-def flatten(seq):
+def flatten(seq: Iterable) -> Iterable:
     """
     Returns a generator object which when evalutated
     returns a flatted version of seq
@@ -119,18 +106,18 @@ def dict_subset(d, keys, include_missing=True, default_key=None):
     return {k: v for k, v in new.items() if v}
 
 
-def get_in_dict(dict_in, keys):
+def get_in_dict(d, keys: Sequence[str]) -> Any:
     """
     Retrieve nested key from dictionary
 
     >>> d = {'a': {'b': {'c': 3}}}
     >>> get_in_dict(d, ('a', 'b', 'c'))
     3
     """
-    return reduce(getitem, keys, dict_in)
+    return reduce(getitem, keys, d)
 
 
-def set_in_dict(dict_in, keys, value):
+def set_in_dict(d: dict, keys: Sequence[str], value) -> None:
     """
     Sets a value inside a nested dictionary
 
@@ -139,7 +126,7 @@ def set_in_dict(dict_in, keys, value):
     >>> d
     {'a': {'b': {'c': 10}}}
     """
-    get_in_dict(dict_in, butlast(keys))[last(keys)] = value
+    get_in_dict(d, butlast(keys))[last(keys)] = value
 
 
 def uniq(seq):
@@ -152,22 +139,22 @@ def uniq(seq):
     return type(seq)(set(seq))
 
 
-def dedupe(items, key=None):
+def dedupe(seq: Sequence, key=None):
     """
     Removes duplicates from a sequence while maintaining order
 
     >>> list(dedupe([1, 5, 2, 1, 9, 1, 5, 10]))
     [1, 5, 2, 9, 10]
     """
-    seen = set()
-    for item in items:
+    seen = set()  # type: set
+    for item in seq:
         val = item if key is None else key(item)
         if val not in seen:
             yield item
             seen.add(val)
 
 
-def prune_dict(d, key=lambda x: x is not None):
+def prune_dict(d: dict, key: Callable = lambda x: x is not None) -> dict:
     """
     Returns new dictionary with key / values pairs filtered by key function.
     Prunes None values by default

requirements.txt

@@ -1,3 +1,5 @@
 coverage
+mypy
+mypy-extensions
 pytest
 sphinx

tests/test_core.py

@@ -5,7 +5,6 @@
 from basic_utils.core import (
     clear,
     getattrs,
-    identity,
     map_getattr,
     recursive_default_dict,
     rgetattr,
@@ -72,14 +71,6 @@ def test_recursive_default_dict():
     assert isinstance(my_dict['random_key'], type(my_dict))
 
 
-@pytest.mark.parametrize("data", [
-    (10), ((10, 20)), (("Homer", "Marge"))
-])
-def test_identity(data):
-    """Tests that identity returns the same arguments as it was passed"""
-    assert identity(data) == data
-
-
 class TestRecursiveGettersAndSetters:
 
     @classmethod

tests/test_helpers.py

@@ -21,7 +21,6 @@
 @pytest.mark.parametrize("data", [
     ([1, 2, 3]),
     ((1, 2, 3,)),
-    ((x for x in [1, 2, 3])),
     (range(1, 4)),
 ])
 def test_first(data):
@@ -32,7 +31,6 @@ def test_first(data):
 @pytest.mark.parametrize("data", [
     ([1, 2, 3]),
     ((1, 2, 3,)),
-    ((x for x in [1, 2, 3])),
     (range(1, 4)),
 ])
 def test_last(data):
@@ -44,7 +42,6 @@ def test_last(data):
     ([1, 2, 3], [2, 3]),
     ((1, 2, 3), (2, 3)),
     ((range(1, 4)), range(2, 4)),
-    ((x for x in [1, 2, 3]), (2, 3)),
 ])
 def test_rest(data, expected):
     """Tests that rest returns the remaining elements in a sequence"""
@@ -54,8 +51,7 @@ def test_rest(data, expected):
 @pytest.mark.parametrize("data, expected", [
     ([1, 2, 3], [1, 2]),
     ((1, 2, 3), (1, 2)),
-    (range(1, 4), range(1, 3)),
-    ((x for x in [1, 2, 3]), (1, 2))
+    (range(1, 4), range(1, 3))
 ])
 def test_butlast(data, expected):
     """Tests that butlast returns all but the last element in a sequence"""
@@ -67,8 +63,6 @@ def test_butlast(data, expected):
     ([1, 2, 3], [3, 2, 1]),
     ((1, 2, 3), (3, 2, 1)),
     (range(1, 4), range(3, 0, -1)),
-    ((x for x in [1, 2, 3]), (3, 2, 1)),
-    ((x for x in []), ())
 ])
 def test_reverse(data, expected):
     assert reverse(data) == expected
@@ -89,13 +83,13 @@ def test_flatten(data, expected):
 @pytest.mark.parametrize("data, expected", [
     ([[1], [2], [3]], [1, 2, 3]),
     ([[1, 2], [3, 4, 5]], [1, 2, 3, 4, 5]),
-    (((1, 2), (3,)), (1, 2, 3)),
+    (((1, 2), (3,)), [1, 2, 3]),
 ])
 def test_partial_flatten(data, expected):
     """
     Tests that partial_flatten returns a partially flattened version data
     """
-    assert partial_flatten(data) == expected
+    assert list(partial_flatten(data)) == expected
 
 
 @pytest.mark.parametrize("keys, expected, default", [