fix typing to go with 3.10 flow not deprecated typing ones. add support for nested ctes, cleanup nested resolver and simplify code there. remove unnecessary guards

Author: collerek

pyproject.toml

@@ -44,6 +44,7 @@ warn_return_any = true
 warn_unused_configs = true
 check_untyped_defs = true
 disallow_untyped_defs = true
+disallow_any_generics = true
 ignore_missing_imports = true
 
 [tool.coverage.run]

sql_metadata/ast_parser.py

@@ -7,8 +7,6 @@
 ``ValueError``).
 """
 
-from typing import Optional
-
 from sqlglot import exp
 from sqlglot.dialects.dialect import DialectType
 
@@ -30,14 +28,14 @@ class ASTParser:
 
     def __init__(self, sql: str) -> None:
         self._raw_sql = sql
-        self._ast: Optional[exp.Expression] = None
+        self._ast: exp.Expression | None = None
         self._dialect: DialectType = None
         self._parsed = False
         self._is_replace = False
         self._cte_name_map: dict[str, str] = {}
 
     @property
-    def ast(self) -> Optional[exp.Expression]:
+    def ast(self) -> exp.Expression | None:
         """The sqlglot AST for the query, lazily parsed on first access.
 
         :returns: Root AST node, or ``None`` for empty/comment-only queries.
@@ -74,7 +72,7 @@ def is_replace(self) -> bool:
         return self._is_replace
 
     @property
-    def cte_name_map(self) -> dict:
+    def cte_name_map(self) -> dict[str, str]:
         """Map of placeholder CTE names back to their original qualified form.
 
         Keys are underscore-separated placeholders (``db__DOT__name``),
@@ -83,7 +81,7 @@ def cte_name_map(self) -> dict:
         _ = self.ast
         return self._cte_name_map
 
-    def _parse(self, sql: str) -> Optional[exp.Expression]:
+    def _parse(self, sql: str) -> exp.Expression | None:
         """Parse *sql* into a sqlglot AST.
 
         Delegates preprocessing to :class:`SqlCleaner` and dialect
@@ -92,7 +90,7 @@ def _parse(self, sql: str) -> Optional[exp.Expression]:
         :param sql: Raw SQL string (may include comments).
         :type sql: str
         :returns: Root AST node, or ``None`` for empty input.
-        :rtype: Optional[exp.Expression]
+        :rtype: exp.Expression | None
         :raises ValueError: If the SQL is malformed.
         """
         if not sql or not sql.strip():

sql_metadata/column_extractor.py

@@ -12,12 +12,12 @@
 """
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Union
+from typing import Any
 
 from sqlglot import exp
 
 from sql_metadata.exceptions import InvalidQueryDefinition
-from sql_metadata.utils import UniqueList, _make_reverse_cte_map, last_segment
+from sql_metadata.utils import UniqueList, last_segment
 
 # ---------------------------------------------------------------------------
 # Result dataclass
@@ -32,13 +32,13 @@ class ExtractionResult:
     """
 
     columns: UniqueList
-    columns_dict: Dict[str, UniqueList]
+    columns_dict: dict[str, UniqueList]
     alias_names: UniqueList
-    alias_dict: Optional[Dict[str, UniqueList]]
-    alias_map: Dict[str, Union[str, list]]
+    alias_dict: dict[str, UniqueList] | None
+    alias_map: dict[str, str | list[str]]
     cte_names: UniqueList
     subquery_names: UniqueList
-    output_columns: list
+    output_columns: list[str]
 
 
 # ---------------------------------------------------------------------------
@@ -47,7 +47,7 @@ class ExtractionResult:
 
 
 #: Simple key → clause-name lookup for most ``arg_types`` keys.
-_CLAUSE_MAP: Dict[str, str] = {
+_CLAUSE_MAP: dict[str, str] = {
     "where": "where",
     "group": "group_by",
     "order": "order_by",
@@ -154,17 +154,17 @@ class _Collector:
         "output_columns",
     )
 
-    def __init__(self, table_aliases: Dict[str, str]):
+    def __init__(self, table_aliases: dict[str, str]):
         self.ta = table_aliases
         self.columns = UniqueList()
-        self.columns_dict: Dict[str, UniqueList] = {}
+        self.columns_dict: dict[str, UniqueList] = {}
         self.alias_names = UniqueList()
-        self.alias_dict: Dict[str, UniqueList] = {}
-        self.alias_map: Dict[str, Union[str, list]] = {}
+        self.alias_dict: dict[str, UniqueList] = {}
+        self.alias_map: dict[str, str | list[str]] = {}
         self.cte_names = UniqueList()
-        self.cte_alias_names: set = set()
-        self.subquery_items: list = []
-        self.output_columns: list = []
+        self.cte_alias_names: set[str] = set()
+        self.subquery_items: list[tuple[int, str]] = []
+        self.output_columns: list[str] = []
 
     def add_column(self, name: str, clause: str) -> None:
         """Record a column name, filing it into the appropriate section."""
@@ -219,14 +219,14 @@ class ColumnExtractor:
     def __init__(
         self,
         ast: exp.Expression,
-        table_aliases: Dict[str, str],
-        cte_name_map: Optional[Dict] = None,
+        table_aliases: dict[str, str],
+        cte_name_map: dict[str, str] | None = None,
     ):
         self._ast = ast
         self._table_aliases = table_aliases
         self._cte_name_map = cte_name_map or {}
         self._collector = _Collector(table_aliases)
-        self._reverse_cte_map = self._build_reverse_cte_map()
+        self._reverse_cte_map = self._cte_name_map
 
     # -------------------------------------------------------------------
     # Public API
@@ -283,19 +283,6 @@ def extract(self) -> ExtractionResult:
     # Setup helpers
     # -------------------------------------------------------------------
 
-    def _build_reverse_cte_map(self) -> Dict[str, str]:
-        """Build a reverse mapping from placeholder CTE names to originals.
-
-        During SQL preprocessing, :class:`SqlCleaner` may rewrite
-        qualified CTE names (e.g. ``schema.cte``) into simple
-        placeholders.  This method inverts that mapping so the final
-        extraction result uses the original qualified names.
-
-        :returns: A dict mapping placeholder names back to their
-            original qualified form.
-        """
-        return _make_reverse_cte_map(self._cte_name_map)
-
     def _seed_cte_names(self) -> None:
         """Pre-populate CTE names in the collector before the main walk.
 
@@ -493,7 +480,9 @@ def _recurse_child(self, child: Any, clause: str, depth: int) -> None:
     # Node handlers
     # -------------------------------------------------------------------
 
-    def _handle_select_exprs(self, exprs: list, clause: str, depth: int) -> None:
+    def _handle_select_exprs(
+        self, exprs: list[exp.Expression], clause: str, depth: int
+    ) -> None:
         """Process the expression list of a SELECT clause.
 
         Iterates each expression in the SELECT list, dispatching to
@@ -920,7 +909,7 @@ def _is_self_alias(
 
     @staticmethod
     def _is_standalone_star(
-        child: exp.Star, seen_stars: set
+        child: exp.Star, seen_stars: set[int]
     ) -> bool:
         """Check whether a star node is standalone (not consumed by a Column).
 
@@ -993,7 +982,7 @@ def _is_cte_with_query_body(
     # Flat column extraction
     # -------------------------------------------------------------------
 
-    def _flat_columns_select_only(self, select: exp.Select) -> list:
+    def _flat_columns_select_only(self, select: exp.Select) -> list[str]:
         """Extract column/alias names from a SELECT's immediate expressions.
 
         Unlike :meth:`_flat_columns`, this does not recurse into the
@@ -1026,7 +1015,7 @@ def _flat_columns_select_only(self, select: exp.Select) -> list:
                     cols.append(col_name)
         return cols
 
-    def _flat_columns(self, node: exp.Expression) -> list:
+    def _flat_columns(self, node: exp.Expression) -> list[str]:
         """Extract all column names from an expression subtree via DFS.
 
         Performs a full depth-first traversal of *node* using
@@ -1052,8 +1041,8 @@ def _flat_columns(self, node: exp.Expression) -> list:
         return cols
 
     def _collect_column_from_node(
-        self, child: exp.Expression, seen_stars: set
-    ) -> Union[str, None]:
+        self, child: exp.Expression, seen_stars: set[int]
+    ) -> str | None:
         """Extract a column name from a single DFS-visited node.
 
         Called by :meth:`_flat_columns` for each node in the traversal.

sql_metadata/comments.py

@@ -19,7 +19,7 @@
 """
 
 import re
-from typing import List
+from typing import Any
 
 from sqlglot.tokens import Tokenizer
 
@@ -71,7 +71,7 @@ def _has_hash_variables(sql: str) -> bool:
     return False
 
 
-def extract_comments(sql: str) -> List[str]:
+def extract_comments(sql: str) -> list[str]:
     """Return all comments found in *sql*, with delimiters preserved.
 
     Tokenizes the SQL, then scans every gap between consecutive token
@@ -107,7 +107,7 @@ def extract_comments(sql: str) -> List[str]:
 _COMMENT_RE = re.compile(r"/\*.*?\*/|/\*.*$|--[^\n]*\n?|#[^\n]*\n?", re.DOTALL)
 
 
-def _scan_gap(sql: str, start: int, end: int, out: list) -> None:
+def _scan_gap(sql: str, start: int, end: int, out: list[str]) -> None:
     """Scan a slice of *sql* for comment delimiters and append matches.
 
     :param sql: The full SQL string (not just the gap).
@@ -118,7 +118,7 @@ def _scan_gap(sql: str, start: int, end: int, out: list) -> None:
     out.extend(_COMMENT_RE.findall(sql[start:end]))
 
 
-def _reconstruct_from_tokens(sql: str, tokens: list) -> str:
+def _reconstruct_from_tokens(sql: str, tokens: list[Any]) -> str:
     """Rebuild SQL from token spans, collapsing gaps to single spaces."""
     if not tokens:
         return ""

sql_metadata/dialect_parser.py

@@ -6,7 +6,7 @@ class so that callers only need to call :meth:`DialectParser.parse`.
 """
 
 import logging
-from typing import Any, Optional
+from typing import Any
 
 import sqlglot
 from sqlglot import exp
@@ -111,7 +111,7 @@ def parse(self, clean_sql: str) -> tuple[exp.Expression, DialectType]:
     # -- dialect detection --------------------------------------------------
 
     @staticmethod
-    def _detect_dialects(sql: str) -> list:
+    def _detect_dialects(sql: str) -> list[Any]:
         """Choose an ordered list of sqlglot dialects to try for *sql*.
 
         Heuristics:
@@ -145,7 +145,7 @@ def _detect_dialects(sql: str) -> list:
     # -- parsing ------------------------------------------------------------
 
     def _try_dialects(
-        self, clean_sql: str, dialects: list
+        self, clean_sql: str, dialects: list[Any]
     ) -> tuple[exp.Expression, DialectType]:
         """Try parsing *clean_sql* with each dialect, returning the best.
 
@@ -180,7 +180,7 @@ def _try_dialects(
         )
 
     @staticmethod
-    def _parse_with_dialect(clean_sql: str, dialect: Any) -> Optional[exp.Expression]:
+    def _parse_with_dialect(clean_sql: str, dialect: Any) -> exp.Expression | None:
         """Parse *clean_sql* with a single dialect, suppressing warnings."""
         logger = logging.getLogger("sqlglot")
         old_level = logger.level