Annotate lazy ops and DSL magic methods

Author: Sergey Vasilyev

Diff for: data_diff/parse_time.py

@@ -33,7 +33,7 @@ class ParseError(ValueError):
 UNITS_STR = ", ".join(sorted(TIME_UNITS.keys()))
 
 
-def string_similarity(a, b):
+def string_similarity(a, b) -> SequenceMatcher:
     return SequenceMatcher(None, a, b).ratio()
 
 

Diff for: data_diff/queries/api.py

@@ -26,31 +26,31 @@ def join(*tables: ITable) -> Join:
     return Join(tables)
 
 
-def leftjoin(*tables: ITable):
+def leftjoin(*tables: ITable) -> Join:
     """Left-joins a sequence of table expressions.
 
     See Also: ``join()``
     """
     return Join(tables, "LEFT")
 
 
-def rightjoin(*tables: ITable):
+def rightjoin(*tables: ITable) -> Join:
     """Right-joins a sequence of table expressions.
 
     See Also: ``join()``
     """
     return Join(tables, "RIGHT")
 
 
-def outerjoin(*tables: ITable):
+def outerjoin(*tables: ITable) -> Join:
     """Outer-joins a sequence of table expressions.
 
     See Also: ``join()``
     """
     return Join(tables, "FULL OUTER")
 
 
-def cte(expr: Expr, *, name: Optional[str] = None, params: Sequence[str] = None):
+def cte(expr: Expr, *, name: Optional[str] = None, params: Sequence[str] = None) -> Cte:
     """Define a CTE"""
     return Cte(expr, name, params)
 
@@ -72,48 +72,48 @@ def table(*path: str, schema: Union[dict, CaseAwareMapping] = None) -> TablePath
     return TablePath(path, schema)
 
 
-def or_(*exprs: Expr):
+def or_(*exprs: Expr) -> Union[BinBoolOp, Expr]:
     """Apply OR between a sequence of boolean expressions"""
     exprs = args_as_tuple(exprs)
     if len(exprs) == 1:
         return exprs[0]
     return BinBoolOp("OR", exprs)
 
 
-def and_(*exprs: Expr):
+def and_(*exprs: Expr) -> Union[BinBoolOp, Expr]:
     """Apply AND between a sequence of boolean expressions"""
     exprs = args_as_tuple(exprs)
     if len(exprs) == 1:
         return exprs[0]
     return BinBoolOp("AND", exprs)
 
 
-def sum_(expr: Expr):
+def sum_(expr: Expr) -> Func:
     """Call SUM(expr)"""
     return Func("sum", [expr])
 
 
-def avg(expr: Expr):
+def avg(expr: Expr) -> Func:
     """Call AVG(expr)"""
     return Func("avg", [expr])
 
 
-def min_(expr: Expr):
+def min_(expr: Expr) -> Func:
     """Call MIN(expr)"""
     return Func("min", [expr])
 
 
-def max_(expr: Expr):
+def max_(expr: Expr) -> Func:
     """Call MAX(expr)"""
     return Func("max", [expr])
 
 
-def exists(expr: Expr):
+def exists(expr: Expr) -> Func:
     """Call EXISTS(expr)"""
     return Func("exists", [expr])
 
 
-def if_(cond: Expr, then: Expr, else_: Optional[Expr] = None):
+def if_(cond: Expr, then: Expr, else_: Optional[Expr] = None) -> CaseWhen:
     """Conditional expression, shortcut to when-then-else.
 
     Example:
@@ -125,7 +125,7 @@ def if_(cond: Expr, then: Expr, else_: Optional[Expr] = None):
     return when(cond).then(then).else_(else_)
 
 
-def when(*when_exprs: Expr):
+def when(*when_exprs: Expr) -> QB_When:
     """Start a when-then expression
 
     Example:
@@ -145,7 +145,7 @@ def when(*when_exprs: Expr):
     return CaseWhen([]).when(*when_exprs)
 
 
-def coalesce(*exprs):
+def coalesce(*exprs) -> Func:
     "Returns a call to COALESCE"
     exprs = args_as_tuple(exprs)
     return Func("COALESCE", exprs)
@@ -160,7 +160,7 @@ def insert_rows_in_batches(db, tbl: TablePath, rows, *, columns=None, batch_size
         db.query(tbl.insert_rows(batch, columns=columns))
 
 
-def current_timestamp():
+def current_timestamp() -> CurrentTimestamp:
     """Returns CURRENT_TIMESTAMP() or NOW()"""
     return CurrentTimestamp()
 

Diff for: data_diff/queries/ast_classes.py

@@ -45,7 +45,7 @@ def _dfs_values(self):
                 if isinstance(v, ExprNode):
                     yield from v._dfs_values()
 
-    def cast_to(self, to):
+    def cast_to(self, to) -> "Cast":
         return Cast(self, to)
 
 
@@ -110,7 +110,7 @@ def select(self, *exprs, distinct=SKIP, optimizer_hints=SKIP, **named_exprs) ->
         resolve_names(self.source_table, exprs)
         return Select.make(self, columns=exprs, distinct=distinct, optimizer_hints=optimizer_hints)
 
-    def where(self, *exprs):
+    def where(self, *exprs) -> "Select":
         """Filter the rows, based on the given predicates. (aka Selection)"""
         exprs = args_as_tuple(exprs)
         exprs = _drop_skips(exprs)
@@ -120,7 +120,7 @@ def where(self, *exprs):
         resolve_names(self.source_table, exprs)
         return Select.make(self, where_exprs=exprs)
 
-    def order_by(self, *exprs):
+    def order_by(self, *exprs) -> "Select":
         """Order the rows lexicographically, according to the given expressions."""
         exprs = _drop_skips(exprs)
         if not exprs:
@@ -129,14 +129,14 @@ def order_by(self, *exprs):
         resolve_names(self.source_table, exprs)
         return Select.make(self, order_by_exprs=exprs)
 
-    def limit(self, limit: int):
+    def limit(self, limit: int) -> "Select":
         """Stop yielding rows after the given limit. i.e. take the first 'n=limit' rows"""
         if limit is SKIP:
             return self
 
         return Select.make(self, limit_expr=limit)
 
-    def join(self, target: "ITable"):
+    def join(self, target: "ITable") -> "Join":
         """Join the current table with the target table, returning a new table containing both side-by-side.
 
         When joining, it's recommended to use explicit tables names, instead of `this`, in order to avoid potential name collisions.
@@ -180,37 +180,37 @@ def group_by(self, *keys) -> "GroupBy":
 
         return GroupBy(self, keys)
 
-    def _get_column(self, name: str):
+    def _get_column(self, name: str) -> "Column":
         if self.schema:
             name = self.schema.get_key(name)  # Get the actual name. Might be case-insensitive.
         return Column(self, name)
 
     # def __getattr__(self, column):
     #     return self._get_column(column)
 
-    def __getitem__(self, column):
+    def __getitem__(self, column) -> "Column":
         if not isinstance(column, str):
             raise TypeError()
         return self._get_column(column)
 
-    def count(self):
+    def count(self) -> "Select":
         """SELECT count() FROM self"""
         return Select(self, [Count()])
 
-    def union(self, other: "ITable"):
+    def union(self, other: "ITable") -> "TableOp":
         """SELECT * FROM self UNION other"""
         return TableOp("UNION", self, other)
 
-    def union_all(self, other: "ITable"):
+    def union_all(self, other: "ITable") -> "TableOp":
         """SELECT * FROM self UNION ALL other"""
         return TableOp("UNION ALL", self, other)
 
-    def minus(self, other: "ITable"):
+    def minus(self, other: "ITable") -> "TableOp":
         """SELECT * FROM self EXCEPT other"""
         # aka
         return TableOp("EXCEPT", self, other)
 
-    def intersect(self, other: "ITable"):
+    def intersect(self, other: "ITable") -> "TableOp":
         """SELECT * FROM self INTERSECT other"""
         return TableOp("INTERSECT", self, other)
 
@@ -233,51 +233,51 @@ def type(self) -> Optional[type]:
 
 @attrs.define(frozen=False, eq=False)
 class LazyOps:
-    def __add__(self, other):
+    def __add__(self, other) -> "BinOp":
         return BinOp("+", [self, other])
 
-    def __sub__(self, other):
+    def __sub__(self, other) -> "BinOp":
         return BinOp("-", [self, other])
 
-    def __neg__(self):
+    def __neg__(self) -> "UnaryOp":
         return UnaryOp("-", self)
 
-    def __gt__(self, other):
+    def __gt__(self, other) -> "BinBoolOp":
         return BinBoolOp(">", [self, other])
 
-    def __ge__(self, other):
+    def __ge__(self, other) -> "BinBoolOp":
         return BinBoolOp(">=", [self, other])
 
-    def __eq__(self, other):
+    def __eq__(self, other) -> "BinBoolOp":
         if other is None:
             return BinBoolOp("IS", [self, None])
         return BinBoolOp("=", [self, other])
 
-    def __lt__(self, other):
+    def __lt__(self, other) -> "BinBoolOp":
         return BinBoolOp("<", [self, other])
 
-    def __le__(self, other):
+    def __le__(self, other) -> "BinBoolOp":
         return BinBoolOp("<=", [self, other])
 
-    def __or__(self, other):
+    def __or__(self, other) -> "BinBoolOp":
         return BinBoolOp("OR", [self, other])
 
-    def __and__(self, other):
+    def __and__(self, other) -> "BinBoolOp":
         return BinBoolOp("AND", [self, other])
 
-    def is_distinct_from(self, other):
+    def is_distinct_from(self, other) -> "IsDistinctFrom":
         return IsDistinctFrom(self, other)
 
-    def like(self, other):
+    def like(self, other) -> "BinBoolOp":
         return BinBoolOp("LIKE", [self, other])
 
-    def sum(self):
+    def sum(self) -> "Func":
         return Func("SUM", [self])
 
-    def max(self):
+    def max(self) -> "Func":
         return Func("MAX", [self])
 
-    def min(self):
+    def min(self) -> "Func":
         return Func("MIN", [self])