Merge pull request #827 from datafold/type-annotate-everything-2

Type annotate some things ("no-brainers")

Author: nolar

data_diff/__main__.py

@@ -77,7 +77,7 @@ def _get_schema(pair: Tuple[Database, DbPath]) -> Dict[str, RawColumnInfo]:
     return db.query_table_schema(table_path)
 
 
-def diff_schemas(table1, table2, schema1, schema2, columns):
+def diff_schemas(table1, table2, schema1, schema2, columns) -> None:
     logging.info("Diffing schemas...")
     attrs = "name", "type", "datetime_precision", "numeric_precision", "numeric_scale"
     for c in columns:
@@ -103,7 +103,7 @@ def diff_schemas(table1, table2, schema1, schema2, columns):
 
 
 class MyHelpFormatter(click.HelpFormatter):
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs) -> None:
         super().__init__(self, **kwargs)
         self.indent_increment = 6
 
@@ -281,7 +281,7 @@ def write_usage(self, prog: str, args: str = "", prefix: Optional[str] = None) -
     default=None,
     help="Override the dbt production schema configuration within dbt_project.yml",
 )
-def main(conf, run, **kw):
+def main(conf, run, **kw) -> None:
     log_handlers = _get_log_handlers(kw["dbt"])
     if kw["table2"] is None and kw["database2"]:
         # Use the "database table table" form
@@ -341,9 +341,7 @@ def main(conf, run, **kw):
                 production_schema_flag=kw["prod_schema"],
             )
         else:
-            return _data_diff(
-                dbt_project_dir=project_dir_override, dbt_profiles_dir=profiles_dir_override, state=state, **kw
-            )
+            _data_diff(dbt_project_dir=project_dir_override, dbt_profiles_dir=profiles_dir_override, state=state, **kw)
     except Exception as e:
         logging.error(e)
         raise
@@ -389,7 +387,7 @@ def _data_diff(
     threads1=None,
     threads2=None,
     __conf__=None,
-):
+) -> None:
     if limit and stats:
         logging.error("Cannot specify a limit when using the -s/--stats switch")
         return

data_diff/abcs/database_types.py

@@ -290,7 +290,7 @@ class Integer(NumericType, IKey):
     precision: int = 0
     python_type: type = int
 
-    def __attrs_post_init__(self):
+    def __attrs_post_init__(self) -> None:
         assert self.precision == 0
 
 

data_diff/cloud/data_source.py

@@ -46,7 +46,7 @@ def process_response(self, value: str) -> str:
         return value
 
 
-def _validate_temp_schema(temp_schema: str):
+def _validate_temp_schema(temp_schema: str) -> None:
     if len(temp_schema.split(".")) != 2:
         raise ValueError("Temporary schema should have a format <database>.<schema>")
 

data_diff/cloud/datafold_api.py

@@ -185,7 +185,7 @@ class DatafoldAPI:
     host: str = "https://app.datafold.com"
     timeout: int = 30
 
-    def __attrs_post_init__(self):
+    def __attrs_post_init__(self) -> None:
         self.host = self.host.rstrip("/")
         self.headers = {
             "Authorization": f"Key {self.api_key}",

data_diff/config.py

@@ -99,7 +99,7 @@ def _apply_config(config: Dict[str, Any], run_name: str, kw: Dict[str, Any]):
 _ENV_VAR_PATTERN = r"\$\{([A-Za-z0-9_]+)\}"
 
 
-def _resolve_env(config: Dict[str, Any]):
+def _resolve_env(config: Dict[str, Any]) -> None:
     """
     Resolve environment variables referenced as ${ENV_VAR_NAME}.
     Missing environment variables are replaced with an empty string.

data_diff/databases/_connect.py

@@ -100,7 +100,7 @@ class Connect:
     database_by_scheme: Dict[str, Database]
     conn_cache: MutableMapping[Hashable, Database]
 
-    def __init__(self, database_by_scheme: Dict[str, Database] = DATABASE_BY_SCHEME):
+    def __init__(self, database_by_scheme: Dict[str, Database] = DATABASE_BY_SCHEME) -> None:
         super().__init__()
         self.database_by_scheme = database_by_scheme
         self.conn_cache = weakref.WeakValueDictionary()

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
@@ -116,7 +131,7 @@ def dialect(self) -> "BaseDialect":
     def compile(self, elem, params=None) -> str:
         return self.dialect.compile(self, elem, params)
 
-    def new_unique_name(self, prefix="tmp"):
+    def new_unique_name(self, prefix="tmp") -> str:
         self._counter[0] += 1
         return f"{prefix}{self._counter[0]}"
 
@@ -173,7 +188,7 @@ class ThreadLocalInterpreter:
     compiler: Compiler
     gen: Generator
 
-    def apply_queries(self, callback: Callable[[str], Any]):
+    def apply_queries(self, callback: Callable[[str], Any]) -> None:
         q: Expr = next(self.gen)
         while True:
             sql = self.compiler.database.dialect.compile(self.compiler, q)
@@ -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):
+    def __len__(self) -> int:
         return len(self.rows)
 
-    def __getitem__(self, i):
+    def __getitem__(self, i) -> Row:
         return self.rows[i]
 
 
@@ -1209,7 +1225,7 @@ class ThreadedDatabase(Database):
     _queue: Optional[ThreadPoolExecutor] = None
     thread_local: threading.local = attrs.field(factory=threading.local)
 
-    def __attrs_post_init__(self):
+    def __attrs_post_init__(self) -> None:
         self._queue = ThreadPoolExecutor(self.thread_count, initializer=self.set_conn)
         logger.info(f"[{self.name}] Starting a threadpool, size={self.thread_count}.")
 

data_diff/databases/bigquery.py

@@ -85,10 +85,10 @@ class Dialect(BaseDialect):
     def random(self) -> str:
         return "RAND()"
 
-    def quote(self, s: str):
+    def quote(self, s: str) -> str:
         return f"`{s}`"
 
-    def to_string(self, s: str):
+    def to_string(self, s: str) -> str:
         return f"cast({s} as string)"
 
     def type_repr(self, t) -> str:
@@ -212,7 +212,7 @@ class BigQuery(Database):
     dataset: str
     _client: Any
 
-    def __init__(self, project, *, dataset, bigquery_credentials=None, **kw):
+    def __init__(self, project, *, dataset, bigquery_credentials=None, **kw) -> None:
         super().__init__()
         credentials = bigquery_credentials
         bigquery = import_bigquery()

data_diff/databases/clickhouse.py

@@ -175,7 +175,7 @@ class Clickhouse(ThreadedDatabase):
 
     _args: Dict[str, Any]
 
-    def __init__(self, *, thread_count: int, **kw):
+    def __init__(self, *, thread_count: int, **kw) -> None:
         super().__init__(thread_count=thread_count)
 
         self._args = kw

data_diff/databases/databricks.py

@@ -65,7 +65,7 @@ def type_repr(self, t) -> str:
         except KeyError:
             return super().type_repr(t)
 
-    def quote(self, s: str):
+    def quote(self, s: str) -> str:
         return f"`{s}`"
 
     def to_string(self, s: str) -> str:
@@ -118,7 +118,7 @@ class Databricks(ThreadedDatabase):
     catalog: str
     _args: Dict[str, Any]
 
-    def __init__(self, *, thread_count, **kw):
+    def __init__(self, *, thread_count, **kw) -> None:
         super().__init__(thread_count=thread_count)
         logging.getLogger("databricks.sql").setLevel(logging.WARNING)