Further refactor

Author: VaggelisD

sqlglot/dialects/snowflake.py

@@ -41,9 +41,6 @@
     from sqlglot._typing import E, B
 
 
-DATE_PARTS = ["YEAR", "QUARTER", "MONTH", "WEEK", "DAY"]
-
-
 def _build_strtok(args: t.List) -> exp.SplitPart:
     # Add default delimiter (space) if missing - per Snowflake docs
     if len(args) == 1:

sqlglot/typing/__init__.py

@@ -311,5 +311,4 @@
     },
     exp.ToMap: {"annotator": lambda self, e: self._annotate_to_map(e)},
     exp.Unnest: {"annotator": lambda self, e: self._annotate_unnest(e)},
-    exp.Unnest: {"annotator": lambda self, e: self._annotate_unnest(e)},
 }

sqlglot/typing/bigquery.py

@@ -192,12 +192,6 @@ def _annotate_array(self: TypeAnnotator, expression: exp.Array) -> exp.Array:
             exp.LaxBool,
         }
     },
-    **{
-        expr_type: {"returns": exp.DataType.Type.DATE}
-        for expr_type in {
-            exp.DateFromUnixDate,
-        }
-    },
     **{
         expr_type: {"returns": exp.DataType.Type.DATETIME}
         for expr_type in {
@@ -272,14 +266,30 @@ def _annotate_array(self: TypeAnnotator, expression: exp.Array) -> exp.Array:
             exp.Uuid,
         }
     },
+    **{
+        expr_type: {"annotator": lambda self, e: _annotate_by_args_with_coerce(self, e)}
+        for expr_type in {
+            exp.PercentileCont,
+            exp.SafeAdd,
+            exp.SafeMultiply,
+            exp.SafeSubtract,
+        }
+    },
+    **{
+        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this", array=True)}
+        for expr_type in {
+            exp.ApproxQuantiles,
+            exp.JSONExtractArray,
+            exp.RegexpExtractAll,
+            exp.Split,
+        }
+    },
     **{
         expr_type: {"returns": exp.DataType.Type.TIMESTAMPTZ} for expr_type in TIMESTAMP_EXPRESSIONS
     },
+    exp.DateFromUnixDate: {"returns": exp.DataType.Type.DATE},
     exp.ParseBignumeric: {"returns": exp.DataType.Type.BIGDECIMAL},
     exp.ParseNumeric: {"returns": exp.DataType.Type.DECIMAL},
-    exp.ApproxQuantiles: {
-        "annotator": lambda self, e: self._annotate_by_args(e, "this", array=True)
-    },
     exp.ApproxTopK: {"annotator": lambda self, e: _annotate_by_args_approx_top(self, e)},
     exp.ApproxTopSum: {"annotator": lambda self, e: _annotate_by_args_approx_top(self, e)},
     exp.Array: {"annotator": _annotate_array},
@@ -292,9 +302,6 @@ def _annotate_array(self: TypeAnnotator, expression: exp.Array) -> exp.Array:
             e, exp.DataType.build("ARRAY<TIMESTAMP>", dialect="bigquery")
         )
     },
-    exp.JSONExtractArray: {
-        "annotator": lambda self, e: self._annotate_by_args(e, "this", array=True)
-    },
     exp.JSONFormat: {
         "annotator": lambda self, e: self._annotate_with_type(
             e, exp.DataType.Type.JSON if e.args.get("to_json") else exp.DataType.Type.VARCHAR
@@ -311,14 +318,6 @@ def _annotate_array(self: TypeAnnotator, expression: exp.Array) -> exp.Array:
         )
     },
     exp.Lag: {"annotator": lambda self, e: self._annotate_by_args(e, "this", "default")},
-    exp.PercentileCont: {"annotator": lambda self, e: _annotate_by_args_with_coerce(self, e)},
-    exp.RegexpExtractAll: {
-        "annotator": lambda self, e: self._annotate_by_args(e, "this", array=True)
-    },
-    exp.SafeAdd: {"annotator": lambda self, e: _annotate_by_args_with_coerce(self, e)},
-    exp.SafeMultiply: {"annotator": lambda self, e: _annotate_by_args_with_coerce(self, e)},
-    exp.SafeSubtract: {"annotator": lambda self, e: _annotate_by_args_with_coerce(self, e)},
-    exp.Split: {"annotator": lambda self, e: self._annotate_by_args(e, "this", array=True)},
     exp.ToCodePoints: {
         "annotator": lambda self, e: self._annotate_with_type(
             e, exp.DataType.build("ARRAY<BIGINT>", dialect="bigquery")

sqlglot/typing/snowflake.py

@@ -36,9 +36,10 @@ def _annotate_timestamp_from_parts(
 
 
 def _annotate_date_or_time_add(self: TypeAnnotator, expression: exp.Expression) -> exp.Expression:
-    DATE_PARTS = {"DAY", "WEEK", "MONTH", "QUARTER", "YEAR"}
     self._annotate_args(expression)
 
+    DATE_PARTS = {"DAY", "WEEK", "MONTH", "QUARTER", "YEAR"}
+
     if (
         expression.this.is_type(exp.DataType.Type.DATE)
         and expression.text("unit").upper() not in DATE_PARTS
@@ -69,9 +70,7 @@ def _annotate_date_or_time_add(self: TypeAnnotator, expression: exp.Expression)
         }
     },
     **{
-        expr_type: {
-            "annotator": lambda self, e: self._annotate_with_type(e, exp.DataType.Type.ARRAY)
-        }
+        expr_type: {"returns": exp.DataType.Type.ARRAY}
         for expr_type in (
             exp.RegexpExtractAll,
             exp.Split,