Annotate regular methods with clear result types

Author: Sergey Vasilyev

data_diff/databases/base.py

@@ -131,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]}"
 

data_diff/databases/bigquery.py

@@ -80,10 +80,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:

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:

data_diff/databases/mssql.py

@@ -76,7 +76,7 @@ class Dialect(BaseDialect):
         "json": JSON,
     }
 
-    def quote(self, s: str):
+    def quote(self, s: str) -> str:
         return f"[{s}]"
 
     def set_timezone_to_utc(self) -> str:
@@ -93,7 +93,7 @@ def current_schema(self) -> str:
         FROM sys.database_principals
         WHERE name = CURRENT_USER"""
 
-    def to_string(self, s: str):
+    def to_string(self, s: str) -> str:
         # Both convert(varchar(max), …) and convert(text, …) do work.
         return f"CONVERT(VARCHAR(MAX), {s})"
 

data_diff/databases/mysql.py

@@ -70,10 +70,10 @@ class Dialect(BaseDialect):
         "boolean": Boolean,
     }
 
-    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 char)"
 
     def is_distinct_from(self, a: str, b: str) -> str:

data_diff/databases/oracle.py

@@ -59,10 +59,10 @@ class Dialect(
     ROUNDS_ON_PREC_LOSS = True
     PLACEHOLDER_TABLE = "DUAL"
 
-    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 varchar(1024))"
 
     def limit_select(

data_diff/databases/vertica.py

@@ -57,7 +57,7 @@ class Dialect(BaseDialect):
         "boolean": Boolean,
     }
 
-    def quote(self, s: str):
+    def quote(self, s: str) -> str:
         return f'"{s}"'
 
     def concat(self, items: List[str]) -> str:

data_diff/parse_time.py

@@ -53,7 +53,7 @@ def parse_time_atom(count, unit):
     return count, unit
 
 
-def parse_time_delta(t: str):
+def parse_time_delta(t: str) -> timedelta:
     time_dict = {}
     while t:
         m = TIME_RE.match(t)
@@ -70,5 +70,5 @@ def parse_time_delta(t: str):
     return timedelta(**time_dict)
 
 
-def parse_time_before(time: datetime, delta: str):
+def parse_time_before(time: datetime, delta: str) -> datetime:
     return time - parse_time_delta(delta)

data_diff/table_segment.py

@@ -138,7 +138,7 @@ def __attrs_post_init__(self) -> None:
                 f"Error: min_update expected to be smaller than max_update! ({self.min_update} >= {self.max_update})"
             )
 
-    def _where(self):
+    def _where(self) -> Optional[str]:
         return f"({self.where})" if self.where else None
 
     def _with_raw_schema(self, raw_schema: Dict[str, RawColumnInfo]) -> Self:

data_diff/tracking.py

@@ -85,7 +85,7 @@ def disable_tracking() -> None:
     g_tracking_enabled = False
 
 
-def is_tracking_enabled():
+def is_tracking_enabled() -> bool:
     return g_tracking_enabled
 
 
@@ -114,7 +114,7 @@ def set_dbt_project_id(s) -> None:
     dbt_project_id = s
 
 
-def get_anonymous_id():
+def get_anonymous_id() -> str:
     global g_anonymous_id
     if g_anonymous_id is None:
         profile = _load_profile()