Merge pull request #20 from kbsriram/main

First batch of type annotations.

Author: FoamyGuy

adafruit_itertools/__init__.py

@@ -28,6 +28,30 @@
 __version__ = "0.0.0+auto.0"
 __repo__ = "https://github.com/adafruit/Adafruit_CircuitPython_Itertools.git"
 
+try:
+ from typing import (
+ Any,
+ Callable,
+ Iterable,
+ Iterator,
+ List,
+ Optional,
+ Sequence,
+ Tuple,
+ TypeVar,
+ Union,
+ )
+ from typing_extensions import TypeAlias
+
+ _T = TypeVar("_T")
+ _Fill = TypeVar("_Fill")
+ _OptionalFill: TypeAlias = Optional[_Fill]
+ _N: TypeAlias = Union[int, float, complex]
+ _Predicate: TypeAlias = Callable[[_T], object]
+
+except ImportError:
+ pass
+
 
 def accumulate(iterable, func=lambda x, y: x + y):
  """Make an iterator that returns accumulated sums, or accumulated
@@ -52,7 +76,7 @@ def accumulate(iterable, func=lambda x, y: x + y):
  yield acc
 
 
-def chain(*iterables):
+def chain(*iterables: Iterable[_T]) -> Iterator[_T]:
  """Make an iterator that returns elements from the first iterable until it
  is exhausted, then proceeds to the next iterable, until all of the iterables
  are exhausted. Used for treating consecutive sequences as a single sequence.
@@ -65,7 +89,7 @@ def chain(*iterables):
  yield from i
 
 
-def chain_from_iterable(iterables):
+def chain_from_iterable(iterables: Iterable[Iterable[_T]]) -> Iterator[_T]:
  """Alternate constructor for chain(). Gets chained inputs from a
  single iterable argument that is evaluated lazily.
 
@@ -78,7 +102,7 @@ def chain_from_iterable(iterables):
  yield element
 
 
-def combinations(iterable, r):
+def combinations(iterable: Iterable[_T], r: int) -> Iterator[Tuple[_T, ...]]:
  """Return r length subsequences of elements from the input iterable.
  Combinations are emitted in lexicographic sort order. So, if the input
  iterable is sorted, the combination tuples will be produced in sorted order.
@@ -113,7 +137,9 @@ def combinations(iterable, r):
  yield tuple(pool[i] for i in indices)
 
 
-def combinations_with_replacement(iterable, r):
+def combinations_with_replacement(
+ iterable: Iterable[_T], r: int
+) -> Iterator[Tuple[_T, ...]]:
  """Return r length subsequences of elements from the input iterable allowing
  individual elements to be repeated more than once.
 
@@ -147,7 +173,7 @@ def combinations_with_replacement(iterable, r):
  yield tuple(pool[i] for i in indices)
 
 
-def compress(data, selectors):
+def compress(data: Iterable[_T], selectors: Iterable[Any]) -> Iterable[_T]:
  """Make an iterator that filters elements from data returning only those
  that have a corresponding element in selectors that evaluates to True.
  Stops when either the data or selectors iterables has been exhausted.
@@ -160,7 +186,7 @@ def compress(data, selectors):
  return (d for d, s in zip(data, selectors) if s)
 
 
-def count(start=0, step=1):
+def count(start: _N = 0, step: _N = 1) -> Iterator[_N]:
  """Make an iterator that returns evenly spaced values starting with number
  start. Often used as an argument to map() to generate consecutive data
  points. Also, used with zip() to add sequence numbers.
@@ -196,7 +222,7 @@ def cycle(p):
  yield from p
 
 
-def dropwhile(predicate, iterable):
+def dropwhile(predicate: _Predicate[_T], iterable: Iterable[_T]) -> Iterator[_T]:
  """Make an iterator that drops elements from the iterable as long as the
  predicate is true; afterwards, returns every element. Note, the iterator
  does not produce any output until the predicate first becomes false, so it
@@ -216,7 +242,7 @@ def dropwhile(predicate, iterable):
  yield x
 
 
-def filterfalse(predicate, iterable):
+def filterfalse(predicate: _Predicate[_T], iterable: Iterable[_T]) -> Iterator[_T]:
  """Make an iterator that filters elements from iterable returning only those
  for which the predicate is False. If predicate is None, return the items
  that are false.
@@ -332,7 +358,9 @@ def islice(p, start, stop=(), step=1):
  return
 
 
-def permutations(iterable, r=None):
+def permutations(
+ iterable: Iterable[_T], r: Optional[int] = None
+) -> Iterator[Tuple[_T, ...]]:
  """Return successive r length permutations of elements in the iterable.
 
  If r is not specified or is None, then r defaults to the length of the
@@ -375,7 +403,7 @@ def permutations(iterable, r=None):
  return
 
 
-def product(*args, r=1):
+def product(*args: Iterable[_T], r: int = 1) -> Iterator[Tuple[_T, ...]]:
  """Cartesian product of input iterables.
 
  Roughly equivalent to nested for-loops in a generator expression. For
@@ -399,14 +427,14 @@ def product(*args, r=1):
  # product('ABCD', 'xy') --> Ax Ay Bx By Cx Cy Dx Dy
  # product(range(2), repeat=3) --> 000 001 010 011 100 101 110 111
  pools = [tuple(pool) for pool in args] * r
- result = [[]]
+ result: List[List[_T]] = [[]]
  for pool in pools:
  result = [x + [y] for x in result for y in pool]
  for prod in result:
  yield tuple(prod)
 
 
-def repeat(el, n=None):
+def repeat(el: _T, n: Optional[int] = None) -> Iterator[_T]:
  """Make an iterator that returns object over and over again. Runs
  indefinitely unless the times argument is specified. Used as argument to
  map() for invariant parameters to the called function. Also used with zip()
@@ -424,7 +452,9 @@ def repeat(el, n=None):
  yield el
 
 
-def starmap(function, iterable):
+def starmap(
+ function: Callable[..., _T], iterable: Iterable[Iterable[Any]]
+) -> Iterator[_T]:
  """Make an iterator that computes the function using arguments obtained from
  the iterable. Used instead of map() when argument parameters are already
  grouped in tuples from a single iterable (the data has been “pre-zipped”).
@@ -439,7 +469,7 @@ def starmap(function, iterable):
  yield function(*args)
 
 
-def takewhile(predicate, iterable):
+def takewhile(predicate: _Predicate[_T], iterable: Iterable[_T]) -> Iterator[_T]:
  """Make an iterator that returns elements from the iterable as long
  as the predicate is true.
 
@@ -455,7 +485,7 @@ def takewhile(predicate, iterable):
  break
 
 
-def tee(iterable, n=2):
+def tee(iterable: Iterable[_T], n: int = 2) -> Sequence[Iterator[_T]]:
  """Return n independent iterators from a single iterable.
 
  :param iterable: the iterator from which to make iterators.
@@ -465,7 +495,9 @@ def tee(iterable, n=2):
  return [iter(iterable) for _ in range(n)]
 
 
-def zip_longest(*args, fillvalue=None):
+def zip_longest(
+ *args: Iterable[_T], fillvalue: _OptionalFill = None
+) -> Iterator[Tuple[Union[_T, _OptionalFill], ...]]:
  """Make an iterator that aggregates elements from each of the
  iterables. If the iterables are of uneven length, missing values are
  filled-in with fillvalue. Iteration continues until the longest
@@ -475,7 +507,7 @@ def zip_longest(*args, fillvalue=None):
  :param fillvalue: value to fill in those missing from shorter iterables
  """
  # zip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
- iterators = [iter(it) for it in args]
+ iterators: List[Iterator[Union[_T, _OptionalFill]]] = [iter(it) for it in args]
  num_active = len(iterators)
  if not num_active:
  return