Annotate simple iterators

Author: Sergey Vasilyev

data_diff/databases/base.py

@@ -5,7 +5,22 @@
 import math
 import sys
 import logging
-from typing import Any, Callable, ClassVar, Dict, Generator, Tuple, Optional, Sequence, Type, List, Union, TypeVar
+from typing import (
+    Any,
+    Callable,
+    ClassVar,
+    Dict,
+    Generator,
+    Iterator,
+    NewType,
+    Tuple,
+    Optional,
+    Sequence,
+    Type,
+    List,
+    Union,
+    TypeVar,
+)
 from functools import partial, wraps
 from concurrent.futures import ThreadPoolExecutor
 import threading
@@ -885,20 +900,21 @@ def optimizer_hints(self, hints: str) -> str:
 
 
 T = TypeVar("T", bound=BaseDialect)
+Row = Sequence[Any]
 
 
 @attrs.define(frozen=True)
 class QueryResult:
-    rows: list
+    rows: List[Row]
     columns: Optional[list] = None
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Row]:
         return iter(self.rows)
 
     def __len__(self) -> int:
         return len(self.rows)
 
-    def __getitem__(self, i):
+    def __getitem__(self, i) -> Row:
         return self.rows[i]
 
 

data_diff/diff_tables.py

@@ -6,7 +6,7 @@
 from enum import Enum
 from contextlib import contextmanager
 from operator import methodcaller
-from typing import Dict, Set, List, Tuple, Iterator, Optional, Union
+from typing import Any, Dict, Set, List, Tuple, Iterator, Optional, Union
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
 import attrs
@@ -89,7 +89,7 @@ class DiffResultWrapper:
     stats: dict
     result_list: list = attrs.field(factory=list)
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Any]:
         yield from self.result_list
         for i in self.diff:
             self.result_list.append(i)

data_diff/schema.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Collection, Iterable, Optional
+from typing import Any, Collection, Iterator, Optional
 
 import attrs
 
@@ -28,7 +28,7 @@ class RawColumnInfo(Collection[Any]):
     collation_name: Optional[str] = None
 
     # It was a tuple once, so we keep it backward compatible temporarily, until remade to classes.
-    def __iter__(self) -> Iterable[Any]:
+    def __iter__(self) -> Iterator[Any]:
         return iter(
             (self.column_name, self.data_type, self.datetime_precision, self.numeric_precision, self.numeric_scale)
         )

data_diff/thread_utils.py

@@ -5,7 +5,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures.thread import _WorkItem
 from time import sleep
-from typing import Callable, Iterator, Optional
+from typing import Any, Callable, Iterator, Optional
 
 import attrs
 
@@ -80,7 +80,7 @@ def _worker(self, fn, *args, **kwargs):
     def submit(self, fn: Callable, *args, priority: int = 0, **kwargs):
         self._futures.append(self._pool.submit(self._worker, fn, *args, priority=priority, **kwargs))
 
-    def __iter__(self) -> Iterator:
+    def __iter__(self) -> Iterator[Any]:
         while True:
             if self._exception:
                 raise self._exception

tests/test_database_types.py

@@ -8,9 +8,11 @@
 import logging
 from decimal import Decimal
 from itertools import islice, repeat, chain
+from typing import Iterator
 
 from parameterized import parameterized
 
+from data_diff.databases.base import Row
 from data_diff.utils import number_to_human
 from data_diff.queries.api import table, commit, this, Code
 from data_diff.queries.api import insert_rows_in_batches
@@ -371,7 +373,7 @@ def __init__(self, table_path, conn) -> None:
         self.table_path = table_path
         self.conn = conn
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Row]:
         last_id = 0
         while True:
             query = (
@@ -402,7 +404,7 @@ def __init__(self, max) -> None:
         super().__init__()
         self.max = max
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[datetime]:
         initial = datetime(2000, 1, 1, 0, 0, 0, 0)
         step = timedelta(seconds=3, microseconds=571)
         return islice(chain(self.MANUAL_FAKES, accumulate(repeat(step), initial=initial)), self.max)
@@ -418,7 +420,7 @@ def __init__(self, max) -> None:
         super().__init__()
         self.max = max
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[int]:
         initial = -128
         step = 1
         return islice(chain(self.MANUAL_FAKES, accumulate(repeat(step), initial=initial)), self.max)
@@ -434,7 +436,7 @@ def __init__(self, max) -> None:
         super().__init__()
         self.max = max
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[bool]:
         return iter(self.MANUAL_FAKES[: self.max])
 
     def __len__(self) -> int:
@@ -465,7 +467,7 @@ def __init__(self, max) -> None:
         super().__init__()
         self.max = max
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[float]:
         initial = -10.0001
         step = 0.00571
         return islice(chain(self.MANUAL_FAKES, accumulate(repeat(step), initial=initial)), self.max)
@@ -482,7 +484,7 @@ def __init__(self, max) -> None:
     def __len__(self) -> int:
         return self.max
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[uuid.UUID]:
         return (uuid.uuid1(i) for i in range(self.max))
 
 
@@ -495,7 +497,7 @@ def __init__(self, max) -> None:
         super().__init__()
         self.max = max
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[str]:
         return iter(self.MANUAL_FAKES[: self.max])
 
     def __len__(self) -> int: