Annotate missing fields

Author: Sergey Vasilyev

data_diff/databases/_connect.py

@@ -95,6 +95,8 @@ def match_path(self, dsn):
 class Connect:
     """Provides methods for connecting to a supported database using a URL or connection dict."""
 
+    database_by_scheme: Dict[str, Database]
+    match_uri_path: Dict[str, MatchUriPath]
     conn_cache: MutableMapping[Hashable, Database]
 
     def __init__(self, database_by_scheme: Dict[str, Database] = DATABASE_BY_SCHEME):

data_diff/databases/base.py

@@ -5,7 +5,7 @@
 import math
 import sys
 import logging
-from typing import Any, Callable, Dict, Generator, Tuple, Optional, Sequence, Type, List, Union, TypeVar
+from typing import Any, Callable, ClassVar, Dict, Generator, Tuple, Optional, Sequence, Type, List, Union, TypeVar
 from functools import partial, wraps
 from concurrent.futures import ThreadPoolExecutor
 import threading
@@ -179,6 +179,9 @@ class ThreadLocalInterpreter:
     Useful for cursor-sensitive operations, such as creating a temporary table.
     """
 
+    compiler: Compiler
+    gen: Generator
+
     def __init__(self, compiler: Compiler, gen: Generator):
         super().__init__()
         self.gen = gen
@@ -238,9 +241,9 @@ def optimizer_hints(self, hints: str) -> str:
 
 
 class BaseDialect(abc.ABC):
-    SUPPORTS_PRIMARY_KEY = False
-    SUPPORTS_INDEXES = False
-    TYPE_CLASSES: Dict[str, type] = {}
+    SUPPORTS_PRIMARY_KEY: ClassVar[bool] = False
+    SUPPORTS_INDEXES: ClassVar[bool] = False
+    TYPE_CLASSES: ClassVar[Dict[str, type]] = {}
     MIXINS = frozenset()
 
     PLACEHOLDER_TABLE = None  # Used for Oracle
@@ -847,14 +850,13 @@ class Database(abc.ABC, _RuntypeHackToFixCicularRefrencedDatabase):
     Instanciated using :meth:`~data_diff.connect`
     """
 
-    default_schema: str = None
-    SUPPORTS_ALPHANUMS = True
-    SUPPORTS_UNIQUE_CONSTAINT = False
-
-    CONNECT_URI_KWPARAMS = []
+    default_schema: ClassVar[str] = None
+    SUPPORTS_ALPHANUMS: ClassVar[bool] = True
+    SUPPORTS_UNIQUE_CONSTAINT: ClassVar[bool] = False
+    CONNECT_URI_KWPARAMS: ClassVar[List[str]] = []
 
-    _interactive = False
-    is_closed = False
+    _interactive: bool = False
+    is_closed: bool = False
 
     @property
     def name(self):
@@ -1121,6 +1123,10 @@ class ThreadedDatabase(Database):
     Used for database connectors that do not support sharing their connection between different threads.
     """
 
+    _init_error: Optional[Exception]
+    _queue: ThreadPoolExecutor
+    thread_local: threading.local
+
     def __init__(self, thread_count=1):
         super().__init__()
         self._init_error = None

data_diff/databases/databricks.py

@@ -1,5 +1,5 @@
 import math
-from typing import Dict, Sequence
+from typing import Any, Dict, Sequence
 import logging
 
 from data_diff.abcs.database_types import (
@@ -104,13 +104,17 @@ class Databricks(ThreadedDatabase):
     CONNECT_URI_HELP = "databricks://:<access_token>@<server_hostname>/<http_path>"
     CONNECT_URI_PARAMS = ["catalog", "schema"]
 
+    default_schema: str
+    catalog: str
+    _args: Dict[str, Any]
+
     def __init__(self, *, thread_count, **kw):
         super().__init__(thread_count=thread_count)
         logging.getLogger("databricks.sql").setLevel(logging.WARNING)
 
         self._args = kw
         self.default_schema = kw.get("schema", "default")
-        self.catalog = self._args.get("catalog", "hive_metastore")
+        self.catalog = kw.get("catalog", "hive_metastore")
 
     def create_connection(self):
         databricks = import_databricks()

data_diff/databases/duckdb.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Any, Dict, Union
 
 from data_diff.utils import match_regexps
 from data_diff.abcs.database_types import (
@@ -139,6 +139,9 @@ class DuckDB(Database):
     CONNECT_URI_HELP = "duckdb://<dbname>@<filepath>"
     CONNECT_URI_PARAMS = ["database", "dbpath"]
 
+    _args: Dict[str, Any]
+    _conn: Any
+
     def __init__(self, **kw):
         super().__init__()
         self._args = kw

data_diff/databases/postgresql.py

@@ -1,5 +1,6 @@
-from typing import List
+from typing import Any, ClassVar, Dict, List, Type
 from data_diff.abcs.database_types import (
+    ColType,
     DbPath,
     JSON,
     Timestamp,
@@ -69,7 +70,7 @@ class PostgresqlDialect(
     SUPPORTS_INDEXES = True
     MIXINS = {Mixin_Schema, Mixin_MD5, Mixin_NormalizeValue, Mixin_RandomSample}
 
-    TYPE_CLASSES = {
+    TYPE_CLASSES: ClassVar[Dict[str, Type[ColType]]] = {
         # Timestamps
         "timestamp with time zone": TimestampTZ,
         "timestamp without time zone": Timestamp,
@@ -125,9 +126,11 @@ class PostgreSQL(ThreadedDatabase):
     SUPPORTS_UNIQUE_CONSTAINT = True
     CONNECT_URI_HELP = "postgresql://<user>:<password>@<host>/<database>"
     CONNECT_URI_PARAMS = ["database?"]
-
     default_schema = "public"
 
+    _args: Dict[str, Any]
+    _conn: Any
+
     def __init__(self, *, thread_count, **kw):
         super().__init__(thread_count=thread_count)
         self._args = kw

data_diff/databases/redshift.py

@@ -1,5 +1,6 @@
-from typing import List, Dict
+from typing import ClassVar, List, Dict, Type
 from data_diff.abcs.database_types import (
+    ColType,
     Float,
     JSON,
     TemporalType,
@@ -53,7 +54,7 @@ def normalize_json(self, value: str, _coltype: JSON) -> str:
 
 class Dialect(PostgresqlDialect, Mixin_MD5, Mixin_NormalizeValue, AbstractMixin_MD5, AbstractMixin_NormalizeValue):
     name = "Redshift"
-    TYPE_CLASSES = {
+    TYPE_CLASSES: ClassVar[Dict[str, Type[ColType]]] = {
         **PostgresqlDialect.TYPE_CLASSES,
         "double": Float,
         "real": Float,

data_diff/thread_utils.py

@@ -45,6 +45,11 @@ class ThreadedYielder(Iterable):
     Priority for the iterator can be provided via the keyword argument 'priority'. (higher runs first)
     """
 
+    _pool: ThreadPoolExecutor
+    _futures: deque
+    _yield: deque
+    _exception: Optional[None]
+
     def __init__(self, max_workers: Optional[int] = None):
         super().__init__()
         self._pool = PriorityThreadPoolExecutor(max_workers)

data_diff/utils.py

@@ -4,7 +4,7 @@
 import re
 import string
 from abc import abstractmethod
-from typing import Any, Dict, Iterable, Iterator, List, MutableMapping, Sequence, TypeVar, Union
+from typing import Any, Dict, Iterable, Iterator, List, MutableMapping, Optional, Sequence, TypeVar, Union
 from urllib.parse import urlparse
 import operator
 import threading
@@ -175,6 +175,9 @@ def alphanums_to_numbers(s1: str, s2: str):
 
 
 class ArithAlphanumeric(ArithString):
+    _str: str
+    _max_len: Optional[int]
+
     def __init__(self, s: str, max_len=None):
         super().__init__()