Fixups

Author: georgesittas

sqlglot/optimizer/annotate_types.py

@@ -53,7 +53,7 @@ def annotate_types(
     Args:
         expression: Expression to annotate.
         schema: Database schema.
-        annotators: Maps expression type to corresponding annotation function.
+        expression_metadata: Maps expression type to corresponding annotation function.
         coerces_to: Maps expression type to set of types that it can be coerced into.
         overwrite_types: Re-annotate the existing AST types.
     Returns:
@@ -64,7 +64,7 @@ def annotate_types(
 
     return TypeAnnotator(
         schema=schema,
-        annotators=annotators,
+        expression_metadata=expression_metadata,
         coerces_to=coerces_to,
         overwrite_types=overwrite_types,
     ).annotate(expression)
@@ -182,7 +182,6 @@ def __init__(
         self,
         schema: Schema,
         expression_metadata: t.Optional[ExpressionMetadataType] = None,
-        annotators: t.Optional[AnnotatorsType] = None,
         coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
         binary_coercions: t.Optional[BinaryCoercions] = None,
         overwrite_types: bool = True,
@@ -213,7 +212,7 @@ def __init__(
         # When set to False, this enables partial annotation by skipping already-annotated nodes
         self._overwrite_types = overwrite_types
 
-    def clean(self) -> None:
+    def clear(self) -> None:
         self._visited.clear()
         self._null_expressions.clear()
         self._setop_column_types.clear()
@@ -236,6 +235,8 @@ def _set_type(
             self._null_expressions.pop(expression_id, None)
 
     def annotate(self, expression: E, annotate_scope: bool = True) -> E:
+        # This flag is used to avoid costly scope traversals when we only care about annotating
+        # non-column expressions (partial type inference), e.g., when simplifying in the optimizer
         if annotate_scope:
             for scope in traverse_scope(expression):
                 self.annotate_scope(scope)

sqlglot/optimizer/simplify.py

@@ -18,11 +18,11 @@
 if t.TYPE_CHECKING:
     from sqlglot.dialects.dialect import DialectType
 
+    DateRange = t.Tuple[datetime.date, datetime.date]
     DateTruncBinaryTransform = t.Callable[
         [exp.Expression, datetime.date, str, Dialect, exp.DataType], t.Optional[exp.Expression]
     ]
 
-DateRange = t.Tuple[datetime.date, datetime.date]
 
 logger = logging.getLogger("sqlglot")
 
@@ -86,7 +86,7 @@ def _func(self, expression: exp.Expression, *args, **kwargs) -> t.Optional[exp.E
             return new_expression
 
         if new_expression != expression:
-            self._annotator.clean()
+            self._annotator.clear()
             new_expression = self._annotator.annotate(
                 expression=new_expression, annotate_scope=False
             )
@@ -493,8 +493,6 @@ def __init__(
 
     CONCATS = (exp.Concat, exp.DPipe)
 
-    DateRange = t.Tuple[datetime.date, datetime.date]
-
     DATETRUNC_BINARY_COMPARISONS: t.Dict[t.Type[exp.Expression], DateTruncBinaryTransform] = {
         exp.LT: lambda l, dt, u, d, t: l
         < date_literal(dt if dt == date_floor(dt, u, d) else date_floor(dt, u, d) + interval(u), t),