Add type stubs across project

rolling/__init__.py

@@ -11,3 +11,27 @@
 from rolling.stats import Mean, Var, Std, Median, Mode, Skew, Kurtosis
 
 __version__ = "0.5.0"
+__all__ = [
+    "Apply",
+    "ApplyPairwise",
+    "Nunique",
+    "Product",
+    "Sum",
+    "Entropy",
+    "PolynomialHash",
+    "All",
+    "Any",
+    "Match",
+    "Min",
+    "Max",
+    "MinHeap",
+    "Monotonic",
+    "JaccardIndex",
+    "Mean",
+    "Var",
+    "Std",
+    "Median",
+    "Mode",
+    "Skew",
+    "Kurtosis",
+]

rolling/apply.py

@@ -1,7 +1,9 @@
 from collections import deque
 from itertools import islice
+from collections.abc import Iterable, Callable
+from typing import Any
 
-from .base import RollingObject
+from .base import RollingObject, WindowType
 
 
 class Apply(RollingObject):
@@ -49,38 +51,38 @@ class Apply(RollingObject):
      [5, 6, 3, 1]]
 
     """
-    def __init__(self, iterable, window_size, window_type="fixed", operation=sum):
+    def __init__(self, iterable: Iterable[Any], window_size: int, window_type: WindowType = "fixed", operation: Callable[..., Any]=sum):
         super().__init__(iterable, window_size, window_type)
         self._operation = operation
 
-    def _init_fixed(self):
+    def _init_fixed(self) -> None:
         self._buffer = deque([None])
         self._buffer.extend(islice(self._iterator, self.window_size - 1))
 
-    def _init_variable(self):
-        self._buffer = deque()
+    def _init_variable(self) -> None:
+        self._buffer: deque[Any] = deque()
 
-    _init_indexed = _init_variable
+    _init_indexed: Callable[..., None] = _init_variable
 
     @property
     def current_value(self):
         return self._operation(self._buffer)
 
-    def _add_new(self, new):
+    def _add_new(self, new) -> None:
         self._buffer.append(new)
 
-    def _remove_old(self):
+    def _remove_old(self) -> None:
         self._buffer.popleft()
 
-    def _update_window(self, new):
+    def _update_window(self, new) -> None:
         self._add_new(new)
         self._remove_old()
 
     @property
-    def _obs(self):
+    def _obs(self) -> int:
         return len(self._buffer)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "Rolling(operation='{}', window_size={}, window_type='{}')".format(
             self._operation.__name__, self.window_size, self.window_type
         )

rolling/apply_pairwise.py

@@ -1,9 +1,9 @@
 from collections import deque
 from itertools import islice
-
+from collections.abc import Iterable, Callable
 from .base_pairwise import RollingPairwise
-
-
+from typing import Any
+from .base import WindowType
 class ApplyPairwise(RollingPairwise):
     """
     Apply a binary function to windows over two iterables.
@@ -38,42 +38,42 @@ class ApplyPairwise(RollingPairwise):
     [1.0, 0.0, -1.0]
 
     """
-    def __init__(self, iterable_1, iterable_2, window_size, function, window_type="fixed"):
+    def __init__(self, iterable_1: Iterable[Any], iterable_2: Iterable[Any], window_size: int, function: Callable[..., Any], window_type: WindowType="fixed") -> None:
         self._buffer_1 = deque(maxlen=window_size)
         self._buffer_2 = deque(maxlen=window_size)
         self._function = function
         super().__init__(iterable_1, iterable_2, window_size=window_size, window_type=window_type)
 
-    def _init_fixed(self):
+    def _init_fixed(self) -> None:
         pairs = zip(self._iterator_1, self._iterator_2)
         for item_1, item_2 in islice(pairs, self.window_size-1):
             self._buffer_1.append(item_1)
             self._buffer_2.append(item_2)
 
-    def _init_variable(self):
+    def _init_variable(self) -> None:
         pass # no action required
 
     @property
     def current_value(self):
         return self._function(self._buffer_1, self._buffer_2)
 
-    def _add_new(self, new_1, new_2):
+    def _add_new(self, new_1, new_2) -> None:
         self._buffer_1.append(new_1)
         self._buffer_2.append(new_2)
 
-    def _remove_old(self):
+    def _remove_old(self) -> None:
         self._buffer_1.popleft()
         self._buffer_2.popleft()
 
-    def _update_window(self, new_1, new_2):
+    def _update_window(self, new_1, new_2) -> None:
         self._buffer_1.append(new_1)
         self._buffer_2.append(new_2)
 
     @property
-    def _obs(self):
+    def _obs(self) -> int:
         return len(self._buffer_1)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "RollingPairwise(operation='{}', window_size={}, window_type='{}')".format(
             self._function.__name__, self.window_size, self.window_type
         )

rolling/arithmetic/__init__.py

@@ -1,3 +1,4 @@
 from .nunique import Nunique
 from .product import Product
 from .sum import Sum
+__all__ = ["Sum", "Product", "Nunique"]

rolling/arithmetic/nunique.py

@@ -1,6 +1,6 @@
 from collections import Counter, deque
 from itertools import islice
-
+from typing import Any
 from rolling.base import RollingObject
 
 
@@ -40,37 +40,37 @@ class Nunique(RollingObject):
 
     """
 
-    def _init_fixed(self):
+    def _init_fixed(self) -> None:
         head = islice(self._iterator, self.window_size - 1)
         self._buffer = deque(head)
         # append a dummy value that is removed when next() is called
         self._buffer.appendleft("dummy_value")
         self._counter = Counter(self._buffer)
 
-    def _init_variable(self):
-        self._buffer = deque()
-        self._counter = Counter()
+    def _init_variable(self) -> None:
+        self._buffer: deque[Any] = deque()
+        self._counter: Counter[int] = Counter()
 
     _init_indexed = _init_variable
 
-    def _update_window(self, new):
+    def _update_window(self, new) -> None:
         # remove oldest value before appending new to buffer
         self._remove_old()
         self._counter[new] += 1
         self._buffer.append(new)
 
-    def _add_new(self, new):
+    def _add_new(self, new) -> None:
         self._counter[new] += 1
         self._buffer.append(new)
 
-    def _remove_old(self):
+    def _remove_old(self) -> None:
         old = self._buffer.popleft()
         self._counter -= Counter([old])
 
     @property
-    def current_value(self):
+    def current_value(self) -> int:
         return len(self._counter)
 
     @property
-    def _obs(self):
+    def _obs(self) -> int:
         return len(self._buffer)

rolling/arithmetic/sum.py

@@ -1,6 +1,6 @@
 from collections import Counter, deque
 from itertools import islice
-
+from typing import Any
 from rolling.base import RollingObject
 
 
@@ -39,33 +39,33 @@ class Sum(RollingObject):
     [13, 11, 15]
 
     """
-    def _init_fixed(self):
+    def _init_fixed(self) -> None:
         head = islice(self._iterator, self.window_size - 1)
         self._buffer = deque(head, maxlen=self.window_size)
         self._buffer.appendleft(0)
-        self._sum = sum(self._buffer)
+        self._sum: int = sum(self._buffer)
 
-    def _init_variable(self):
-        self._buffer = deque()
+    def _init_variable(self) -> None:
+        self._buffer: deque[int] = deque()
         self._sum = 0
 
     _init_indexed = _init_variable
 
-    def _update_window(self, new):
+    def _update_window(self, new) -> None:
         self._sum += new - self._buffer.popleft()
         self._buffer.append(new)
 
-    def _add_new(self, new):
+    def _add_new(self, new) -> None:
         self._sum += new
         self._buffer.append(new)
 
     def _remove_old(self):
         self._sum -= self._buffer.popleft()
 
     @property
-    def current_value(self):
+    def current_value(self) -> int:
         return self._sum
 
     @property
-    def _obs(self):
+    def _obs(self) -> int:
         return len(self._buffer)

rolling/base.py

@@ -1,10 +1,12 @@
 import abc
 from collections import deque
-from collections.abc import Iterator
+from collections.abc import Iterator, Iterable
 from itertools import chain
+from typing import Any, Literal
 
-
-class RollingObject(Iterator):
+WindowType = Literal["fixed", "variable", "indexed"]
+Tracker = Literal["sortedlist", "skiplist"]
+class RollingObject(Iterator[Any]):
     """
     Baseclass for rolling iterator objects.
 
@@ -31,10 +33,10 @@ class RollingObject(Iterator):
 
     """
 
-    def __init__(self, iterable, window_size, window_type="fixed"):
+    def __init__(self, iterable: Iterable[Any], window_size: int, window_type: WindowType="fixed"):
         self.window_type = window_type
         self.window_size = _validate_window_size(window_size, window_type)
-        self._iterator = iter(iterable)
+        self._iterator: Iterator[Any] = iter(iterable)
         self._filled = self.window_type == "fixed"
 
         if window_type == "fixed":
@@ -47,12 +49,9 @@ def __init__(self, iterable, window_size, window_type="fixed"):
             # Keep track of all indexes that we encounter. Assumes that all
             # values we encounter will be stored in the same order. If not,
             # the subtype will need to implement its own _next_indexed() method.
-            self.index_buffer = deque()
+            self.index_buffer: deque[Any] = deque()
             self._init_indexed()
 
-        else:
-            raise ValueError(f"Unknown window_type '{window_type}'")
-
     def __repr__(self):
         return "Rolling(operation='{}', window_size={}, window_type='{}')".format(
             self.__class__.__name__, self.window_size, self.window_type
@@ -116,7 +115,7 @@ def __next__(self):
 
         raise NotImplementedError(f"next() not implemented for {self.window_type}")
 
-    def extend(self, iterable):
+    def extend(self, iterable: Iterable[Any]):
         """
         Extend the iterator being consumed with a new iterable.
 
@@ -196,13 +195,11 @@ def _update_window(self, new):
         pass
 
 
-def _validate_window_size(window_size, window_type):
+def _validate_window_size(window_size: int, window_type: WindowType) -> int:
     """
     Check if k is a positive integer
     """
     if window_type in {"fixed", "variable"}:
-        if not isinstance(window_size, int):
-            raise TypeError(f"window_size must be integer type, got {type(window_size).__name__}")
         if window_size <= 0:
             raise ValueError("window_size must be positive")
     return window_size

rolling/base_pairwise.py

@@ -1,14 +1,16 @@
 import abc
+from typing import Any
 from collections.abc import Iterator
 from itertools import chain
+from collections.abc import Iterable
+from .base import WindowType
 
-
-class RollingPairwise(Iterator):
+class RollingPairwise(Iterator[Any]):
     """
     Baseclass for rolling iterators over two iterables.
 
     """
-    def __init__(self, iterable_1, iterable_2, window_size, window_type="fixed"):
+    def __init__(self, iterable_1: Iterable[Any], iterable_2: Iterable[Any], window_size: int, window_type: WindowType="fixed"):
         self.window_type = window_type
         self.window_size = _validate_window_size(window_size)
         self._iterator_1 = iter(iterable_1)
@@ -118,12 +120,10 @@ def _update_window(self, new_1, new_2):
         pass
 
 
-def _validate_window_size(k):
+def _validate_window_size(k: int):
     """
     Check if k is a positive integer
     """
-    if not isinstance(k, int):
-        raise TypeError(f"window_size must be integer type, got {type(k).__name__}")
     if k <= 0:
         raise ValueError("window_size must be positive")
     return k