Improve Type Annotations Coverage (#162)

* Add some type annotations

* More type annotation.

* Freedom from typlessness.

* More annotations.

* Clean up linter issues

dbt_common/context.py

@@ -2,14 +2,15 @@
 from typing import List, Mapping, Optional
 
 from dbt_common.constants import PRIVATE_ENV_PREFIX, SECRET_ENV_PREFIX
+from dbt_common.record import Recorder
 
 
 class InvocationContext:
     def __init__(self, env: Mapping[str, str]):
         self._env = {k: v for k, v in env.items() if not k.startswith(PRIVATE_ENV_PREFIX)}
         self._env_secrets: Optional[List[str]] = None
         self._env_private = {k: v for k, v in env.items() if k.startswith(PRIVATE_ENV_PREFIX)}
-        self.recorder = None
+        self.recorder: Optional[Recorder] = None
         # This class will also eventually manage the invocation_id, flags, event manager, etc.
 
     @property
@@ -32,7 +33,7 @@ def env_secrets(self) -> List[str]:
 _INVOCATION_CONTEXT_VAR: ContextVar[InvocationContext] = ContextVar("DBT_INVOCATION_CONTEXT_VAR")
 
 
-def reliably_get_invocation_var() -> ContextVar:
+def reliably_get_invocation_var() -> ContextVar[InvocationContext]:
     invocation_var: Optional[ContextVar] = next(
         (cv for cv in copy_context() if cv.name == _INVOCATION_CONTEXT_VAR.name), None
     )

dbt_common/dataclass_schema.py

@@ -1,4 +1,4 @@
-from typing import ClassVar, cast, get_type_hints, List, Tuple, Dict, Any, Optional
+from typing import Any, cast, ClassVar, Dict, get_type_hints, List, Optional, Tuple
 import re
 import jsonschema
 from dataclasses import fields, Field
@@ -26,7 +26,7 @@ class ValidationError(jsonschema.ValidationError):
 
 
 class DateTimeSerialization(SerializationStrategy):
-    def serialize(self, value) -> str:
+    def serialize(self, value: datetime) -> str:
         out = value.isoformat()
         # Assume UTC if timezone is missing
         if value.tzinfo is None:
@@ -127,7 +127,7 @@ def _get_fields(cls) -> List[Tuple[Field, str]]:
 
     # copied from hologram. Used in tests
     @classmethod
-    def _get_field_names(cls):
+    def _get_field_names(cls) -> List[str]:
         return [element[1] for element in cls._get_fields()]
 
 
@@ -152,7 +152,7 @@ def validate(cls, value):
 
 # These classes must be in this order or it doesn't work
 class StrEnum(str, SerializableType, Enum):
-    def __str__(self):
+    def __str__(self) -> str:
         return self.value
 
     # https://docs.python.org/3.6/library/enum.html#using-automatic-values

dbt_common/exceptions/base.py

@@ -1,5 +1,5 @@
 import builtins
-from typing import List, Any, Optional
+from typing import Any, List, Optional
 import os
 
 from dbt_common.constants import SECRET_ENV_PREFIX
@@ -37,7 +37,7 @@ def __init__(self, msg: str):
         self.msg = scrub_secrets(msg, env_secrets())
 
     @property
-    def type(self):
+    def type(self) -> str:
         return "Internal"
 
     def process_stack(self):
@@ -59,7 +59,7 @@ def process_stack(self):
 
         return lines
 
-    def __str__(self):
+    def __str__(self) -> str:
         if hasattr(self.msg, "split"):
             split_msg = self.msg.split("\n")
         else:

dbt_common/helper_types.py

@@ -19,7 +19,7 @@
 class NVEnum(StrEnum):
     novalue = "novalue"
 
-    def __eq__(self, other):
+    def __eq__(self, other) -> bool:
         return isinstance(other, NVEnum)
 
 
@@ -59,7 +59,7 @@ def includes(self, item_name: str) -> bool:
             item_name in self.include or self.include in self.INCLUDE_ALL
         ) and item_name not in self.exclude
 
-    def _validate_items(self, items: List[str]):
+    def _validate_items(self, items: List[str]) -> None:
         pass
 
 

dbt_common/record.py

@@ -14,8 +14,6 @@
 from enum import Enum
 from typing import Any, Callable, Dict, List, Mapping, Optional, Type
 
-from dbt_common.context import get_invocation_context
-
 
 class Record:
     """An instance of this abstract Record class represents a request made by dbt
@@ -295,9 +293,11 @@ def record_function_inner(func_to_record):
             return func_to_record
 
         @functools.wraps(func_to_record)
-        def record_replay_wrapper(*args, **kwargs):
-            recorder: Recorder = None
+        def record_replay_wrapper(*args, **kwargs) -> Any:
+            recorder: Optional[Recorder] = None
             try:
+                from dbt_common.context import get_invocation_context
+
                 recorder = get_invocation_context().recorder
             except LookupError:
                 pass

dbt_common/semver.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 import re
-from typing import List
+from typing import List, Iterable
 
 import dbt_common.exceptions.base
 from dbt_common.exceptions import VersionsNotCompatibleError
@@ -74,7 +74,7 @@ def _cmp(a, b):
 
 @dataclass
 class VersionSpecifier(VersionSpecification):
-    def to_version_string(self, skip_matcher=False):
+    def to_version_string(self, skip_matcher: bool = False) -> str:
         prerelease = ""
         build = ""
         matcher = ""
@@ -92,7 +92,7 @@ def to_version_string(self, skip_matcher=False):
         )
 
     @classmethod
-    def from_version_string(cls, version_string):
+    def from_version_string(cls, version_string: str) -> "VersionSpecifier":
         match = _VERSION_REGEX.match(version_string)
 
         if not match:
@@ -104,7 +104,7 @@ def from_version_string(cls, version_string):
 
         return cls.from_dict(matched)
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.to_version_string()
 
     def to_range(self) -> "VersionRange":
@@ -192,32 +192,32 @@ def compare(self, other):
 
         return 0
 
-    def __lt__(self, other):
+    def __lt__(self, other) -> bool:
         return self.compare(other) == -1
 
-    def __gt__(self, other):
+    def __gt__(self, other) -> bool:
         return self.compare(other) == 1
 
-    def __eq___(self, other):
+    def __eq___(self, other) -> bool:
         return self.compare(other) == 0
 
     def __cmp___(self, other):
         return self.compare(other)
 
     @property
-    def is_unbounded(self):
+    def is_unbounded(self) -> bool:
         return False
 
     @property
-    def is_lower_bound(self):
+    def is_lower_bound(self) -> bool:
         return self.matcher in [Matchers.GREATER_THAN, Matchers.GREATER_THAN_OR_EQUAL]
 
     @property
-    def is_upper_bound(self):
+    def is_upper_bound(self) -> bool:
         return self.matcher in [Matchers.LESS_THAN, Matchers.LESS_THAN_OR_EQUAL]
 
     @property
-    def is_exact(self):
+    def is_exact(self) -> bool:
         return self.matcher == Matchers.EXACT
 
     @classmethod
@@ -418,7 +418,7 @@ def reduce_versions(*args):
     return to_return
 
 
-def versions_compatible(*args):
+def versions_compatible(*args) -> bool:
     if len(args) == 1:
         return True
 
@@ -429,7 +429,7 @@ def versions_compatible(*args):
         return False
 
 
-def find_possible_versions(requested_range, available_versions):
+def find_possible_versions(requested_range, available_versions: Iterable[str]):
     possible_versions = []
 
     for version_string in available_versions:
@@ -442,7 +442,9 @@ def find_possible_versions(requested_range, available_versions):
     return [v.to_version_string(skip_matcher=True) for v in sorted_versions]
 
 
-def resolve_to_specific_version(requested_range, available_versions):
+def resolve_to_specific_version(
+    requested_range, available_versions: Iterable[str]
+) -> Optional[str]:
     max_version = None
     max_version_string = None
 

dbt_common/utils/casting.py

@@ -1,7 +1,7 @@
 # This is useful for proto generated classes in particular, since
 # the default for protobuf for strings is the empty string, so
 # Optional[str] types don't work for generated Python classes.
-from typing import Optional
+from typing import Any, Dict, Optional
 
 
 def cast_to_str(string: Optional[str]) -> str:
@@ -18,8 +18,8 @@ def cast_to_int(integer: Optional[int]) -> int:
         return integer
 
 
-def cast_dict_to_dict_of_strings(dct):
-    new_dct = {}
+def cast_dict_to_dict_of_strings(dct: Dict[Any, Any]) -> Dict[str, str]:
+    new_dct: Dict[str, str] = {}
     for k, v in dct.items():
         new_dct[str(k)] = str(v)
     return new_dct

dbt_common/utils/executor.py

@@ -1,9 +1,12 @@
 import concurrent.futures
 from contextlib import contextmanager
-from contextvars import ContextVar
 from typing import Protocol, Optional
 
-from dbt_common.context import get_invocation_context, reliably_get_invocation_var
+from dbt_common.context import (
+    get_invocation_context,
+    reliably_get_invocation_var,
+    InvocationContext,
+)
 
 
 class ConnectingExecutor(concurrent.futures.Executor):
@@ -63,7 +66,7 @@ class HasThreadingConfig(Protocol):
     threads: Optional[int]
 
 
-def _thread_initializer(invocation_context: ContextVar) -> None:
+def _thread_initializer(invocation_context: InvocationContext) -> None:
     invocation_var = reliably_get_invocation_var()
     invocation_var.set(invocation_context)
 

dbt_common/utils/jinja.py

@@ -5,19 +5,21 @@
 DOCS_PREFIX = "dbt_docs__"
 
 
-def get_dbt_macro_name(name):
+def get_dbt_macro_name(name) -> str:
     if name is None:
         raise DbtInternalError("Got None for a macro name!")
     return f"{MACRO_PREFIX}{name}"
 
 
-def get_dbt_docs_name(name):
+def get_dbt_docs_name(name) -> str:
     if name is None:
         raise DbtInternalError("Got None for a doc name!")
     return f"{DOCS_PREFIX}{name}"
 
 
-def get_materialization_macro_name(materialization_name, adapter_type=None, with_prefix=True):
+def get_materialization_macro_name(
+    materialization_name, adapter_type=None, with_prefix=True
+) -> str:
     if adapter_type is None:
         adapter_type = "default"
     name = f"materialization_{materialization_name}_{adapter_type}"

tests/unit/test_agate_helper.py

@@ -46,13 +46,13 @@
 
 
 class TestAgateHelper(unittest.TestCase):
-    def setUp(self):
+    def setUp(self) -> None:
         self.tempdir = mkdtemp()
 
-    def tearDown(self):
+    def tearDown(self) -> None:
         rmtree(self.tempdir)
 
-    def test_from_csv(self):
+    def test_from_csv(self) -> None:
         path = os.path.join(self.tempdir, "input.csv")
         with open(path, "wb") as fp:
             fp.write(SAMPLE_CSV_DATA.encode("utf-8"))
@@ -61,7 +61,7 @@ def test_from_csv(self):
         for idx, row in enumerate(tbl):
             self.assertEqual(list(row), EXPECTED[idx])
 
-    def test_bom_from_csv(self):
+    def test_bom_from_csv(self) -> None:
         path = os.path.join(self.tempdir, "input.csv")
         with open(path, "wb") as fp:
             fp.write(SAMPLE_CSV_BOM_DATA.encode("utf-8"))
@@ -70,7 +70,7 @@ def test_bom_from_csv(self):
         for idx, row in enumerate(tbl):
             self.assertEqual(list(row), EXPECTED[idx])
 
-    def test_from_csv_all_reserved(self):
+    def test_from_csv_all_reserved(self) -> None:
         path = os.path.join(self.tempdir, "input.csv")
         with open(path, "wb") as fp:
             fp.write(SAMPLE_CSV_DATA.encode("utf-8"))
@@ -79,7 +79,7 @@ def test_from_csv_all_reserved(self):
         for expected, row in zip(EXPECTED_STRINGS, tbl):
             self.assertEqual(list(row), expected)
 
-    def test_from_data(self):
+    def test_from_data(self) -> None:
         column_names = ["a", "b", "c", "d", "e", "f", "g"]
         data = [
             {
@@ -106,7 +106,7 @@ def test_from_data(self):
         for idx, row in enumerate(tbl):
             self.assertEqual(list(row), EXPECTED[idx])
 
-    def test_datetime_formats(self):
+    def test_datetime_formats(self) -> None:
         path = os.path.join(self.tempdir, "input.csv")
         datetimes = [
             "20180806T11:33:29.000Z",
@@ -120,7 +120,7 @@ def test_datetime_formats(self):
             tbl = agate_helper.from_csv(path, ())
             self.assertEqual(tbl[0][0], expected)
 
-    def test_merge_allnull(self):
+    def test_merge_allnull(self) -> None:
         t1 = agate_helper.table_from_rows([(1, "a", None), (2, "b", None)], ("a", "b", "c"))
         t2 = agate_helper.table_from_rows([(3, "c", None), (4, "d", None)], ("a", "b", "c"))
         result = agate_helper.merge_tables([t1, t2])
@@ -130,7 +130,7 @@ def test_merge_allnull(self):
         assert isinstance(result.column_types[2], agate_helper.Integer)
         self.assertEqual(len(result), 4)
 
-    def test_merge_mixed(self):
+    def test_merge_mixed(self) -> None:
         t1 = agate_helper.table_from_rows(
             [(1, "a", None, None), (2, "b", None, None)], ("a", "b", "c", "d")
         )
@@ -181,7 +181,7 @@ def test_merge_mixed(self):
         assert isinstance(result.column_types[3], agate.data_types.Number)
         self.assertEqual(len(result), 6)
 
-    def test_nocast_string_types(self):
+    def test_nocast_string_types(self) -> None:
         # String fields should not be coerced into a representative type
         # See: https://github.com/dbt-labs/dbt-core/issues/2984
 
@@ -202,7 +202,7 @@ def test_nocast_string_types(self):
         for i, row in enumerate(tbl):
             self.assertEqual(list(row), expected[i])
 
-    def test_nocast_bool_01(self):
+    def test_nocast_bool_01(self) -> None:
         # True and False values should not be cast to 1 and 0, and vice versa
         # See: https://github.com/dbt-labs/dbt-core/issues/4511