recce-nightly 1.2.0.20250506__py3-none-any.whl → 1.26.0.20251124__py3-none-any.whl

recce/util/cll.py

@@ -1,19 +1,24 @@
 import time
 from dataclasses import dataclass
-from typing import Dict, List, Literal
+from typing import Dict, List, Optional, Tuple
 
-from sqlglot import parse_one, Dialect
-from sqlglot.errors import SqlglotError, OptimizeError
-from sqlglot.expressions import Column, Alias, Func, Binary, Paren, Case, Expression, If, Union, Intersect
-from sqlglot.optimizer import traverse_scope, Scope
+import sqlglot.expressions as exp
+from sqlglot import Dialect, parse_one
+from sqlglot.errors import OptimizeError, SqlglotError
+from sqlglot.optimizer import Scope, traverse_scope
 from sqlglot.optimizer.qualify import qualify
 
 from recce.exceptions import RecceException
-from recce.util import SingletonMeta
+from recce.models.types import CllColumn, CllColumnDep
+
+CllResult = Tuple[
+    List[CllColumnDep],  # Model to column dependencies
+    Dict[str, CllColumn],  # Column to column dependencies
+]
 
 
 @dataclass
-class CLLPerformanceTracking(metaclass=SingletonMeta):
+class CLLPerformanceTracking:
     lineage_start = None
     lineage_elapsed = None
     column_lineage_start = None
@@ -50,11 +55,11 @@ class CLLPerformanceTracking(metaclass=SingletonMeta):
 
     def to_dict(self):
         return {
-            'lineage_elapsed_ms': self.lineage_elapsed,
-            'column_lineage_elapsed_ms': self.column_lineage_elapsed,
-            'total_nodes': self.total_nodes,
-            'sqlglot_error_nodes': self.sqlglot_error_nodes,
-            'other_error_nodes': self.other_error_nodes
+            "lineage_elapsed_ms": self.lineage_elapsed,
+            "column_lineage_elapsed_ms": self.column_lineage_elapsed,
+            "total_nodes": self.total_nodes,
+            "sqlglot_error_nodes": self.sqlglot_error_nodes,
+            "other_error_nodes": self.other_error_nodes,
         }
 
     def reset(self):
@@ -68,233 +73,283 @@ class CLLPerformanceTracking(metaclass=SingletonMeta):
         self.other_error_nodes = 0
 
 
-@dataclass
-class ColumnLevelDependsOn:
-    node: str
-    column: str
-
-
-@dataclass
-class ColumnLevelDependencyColumn:
-    type: Literal['source', 'passthrough', 'renamed', 'derived']
-    depends_on: List[ColumnLevelDependsOn]
-
-
-def _cll_expression(expression, table_alias_map) -> ColumnLevelDependencyColumn:
+def _cll_column(proj, table_alias_map) -> CllColumn:
     # given an expression, return the columns depends on
     # [{node: table, column: column}, ...]
+    type = "source"
+    depends_on: List[CllColumnDep] = []
+
+    # instance of Column
+    if isinstance(proj, exp.Alias):
+        # 'select a as b'
+        # 'select CURRENT_TIMESTAMP() as create_at'
+        root = proj.this
+
+    for expression in root.walk(bfs=False):
+        if isinstance(expression, exp.Column):
+            column = expression
+            alias = column.table
+
+            if alias is None:
+                table = next(iter(table_alias_map.values()))
+            else:
+                table = table_alias_map.get(alias, alias)
+            depends_on.append(CllColumnDep(table, column.name))
+            if type == "source":
+                type = "passthrough"
+        elif isinstance(expression, (exp.Paren, exp.Identifier)):
+            pass
+        else:
+            type = "derived"
+
+    depends_on = _dedeup_depends_on(depends_on)
+
+    if len(depends_on) == 0:
+        type = "source"
+
+    if isinstance(proj, exp.Alias):
+        alias = proj
+        if type == "passthrough" and depends_on[0].column != alias.alias_or_name:
+            type = "renamed"
+
+    return CllColumn(type=type, depends_on=depends_on)
+
+
+def _dedeup_depends_on(depends_on: List[CllColumnDep]) -> List[CllColumnDep]:
+    # deduplicate the depends_on list
+    dedup_set = set()
+    dedup_list = []
+    for col_dep in depends_on:
+        node_col = col_dep.node + "." + col_dep.column
+        if node_col not in dedup_set:
+            dedup_list.append(col_dep)
+            dedup_set.add(node_col)
+    return dedup_list
+
+
+def _cll_set_scope(scope: Scope, scope_cll_map: dict[Scope, CllResult]) -> CllResult:
+    # model-to-column
+    m2c: List[CllColumnDep] = []
+    # column-to-column
+    c2c_map: Dict[str, CllColumn] = {}
+
+    for union_scope in scope.union_scopes:
+        sub_scope_result = scope_cll_map.get(union_scope)
+        if sub_scope_result is None:
+            raise RecceException(f"Scope {union_scope} not found in scope_cll_map")
+        sub_m2c, sub_c2c_map = sub_scope_result
+
+        for k, v in sub_c2c_map.items():
+            if k not in c2c_map:
+                c2c_map[k] = v
+            else:
+                c2c_map[k].depends_on.extend(v.depends_on)
+                c2c_map[k].transformation_type = "derived"
+
+        m2c.extend(sub_m2c)
+    return m2c, c2c_map
+
+
+def _cll_select_scope(scope: Scope, scope_cll_map: dict[Scope, CllResult]) -> CllResult:
+    assert scope.expression.key == "select"
+
+    # model-to-column
+    m2c: List[CllColumnDep] = []
+    # column-to-column
+    c2c_map: Dict[str, CllColumn] = {}
+
+    table_alias_map = {t.alias_or_name: t.name for t in scope.tables}
+    select = scope.expression
+
+    def source_column_dependency(ref_column: exp.Column) -> Optional[CllColumn]:
+        column_name = ref_column.name
+        table_name = ref_column.table if ref_column.table != "" else next(iter(table_alias_map.values()))
+        source = scope.sources.get(table_name, None)  # transformation_type: exp.Table | Scope
+        if isinstance(source, Scope):
+            ref_cll_result = scope_cll_map.get(source)
+            if ref_cll_result is None:
+                return None
+            _, sub_c2c_map = ref_cll_result
+            return sub_c2c_map.get(column_name)
+        elif isinstance(source, exp.Table):
+            return CllColumn(
+                name=column_name,
+                transformation_type="passthrough",
+                depends_on=[CllColumnDep(node=source.name, column=column_name)],
+            )
+        else:
+            return None
+
+    def subquery_cll(subquery: exp.Subquery) -> Optional[CllResult]:
+        select = subquery.find(exp.Select)
+        if select is None:
+            return None
+
+        matched_scope = None
+        for sub_scope in scope.subquery_scopes:
+            if sub_scope.expression == select:
+                matched_scope = sub_scope
+                break
+        if matched_scope is None:
+            return None
+
+        return scope_cll_map.get(matched_scope)
+
+    for proj in scope.expression.selects:
+        transformation_type = "source"
+        column_depends_on: List[CllColumnDep] = []
+        root = proj.this if isinstance(proj, exp.Alias) else proj
+        for expression in root.walk(bfs=False):
+            if isinstance(expression, exp.Column):
+                ref_column_dependency = source_column_dependency(expression)
+                if ref_column_dependency is not None:
+                    column_depends_on.extend(ref_column_dependency.depends_on)
+                    if ref_column_dependency.transformation_type == "derived":
+                        transformation_type = "derived"
+                    elif ref_column_dependency.transformation_type == "renamed":
+                        if transformation_type == "source" or transformation_type == "passthrough":
+                            transformation_type = "renamed"
+                    elif ref_column_dependency.transformation_type == "passthrough":
+                        if transformation_type == "source":
+                            transformation_type = "passthrough"
+                else:
+                    column_depends_on.append(CllColumnDep(expression.table, expression.name))
+                    if transformation_type == "source":
+                        transformation_type = "passthrough"
 
-    if isinstance(expression, Column):
-        column = expression
-        alias = column.table
+            elif isinstance(expression, (exp.Paren, exp.Identifier)):
+                pass
+            else:
+                transformation_type = "derived"
 
-        if alias is None:
-            table = next(iter(table_alias_map.values()))
-        else:
-            table = table_alias_map.get(alias, alias)
+        column_depends_on = _dedeup_depends_on(column_depends_on)
+
+        if len(column_depends_on) == 0 and transformation_type != "source":
+            transformation_type = "source"
+
+        if isinstance(proj, exp.Alias):
+            alias = proj
+            if transformation_type == "passthrough" and column_depends_on[0].column != alias.alias_or_name:
+                transformation_type = "renamed"
 
-        return ColumnLevelDependencyColumn(
-            type='passthrough',
-            depends_on=[ColumnLevelDependsOn(table, column.name)]
+        c2c_map[proj.alias_or_name] = CllColumn(
+            name=proj.alias_or_name, transformation_type=transformation_type, depends_on=column_depends_on
         )
-    elif isinstance(expression, Paren):
-        return _cll_expression(expression.this, table_alias_map)
-    elif isinstance(expression, Binary):
-        depends_on = []
-        if expression.left:
-            depends_on_left = _cll_expression(expression.left, table_alias_map).depends_on
-            depends_on.extend(depends_on_left)
-        if expression.right:
-            depends_on_right = _cll_expression(expression.right, table_alias_map).depends_on
-            depends_on.extend(depends_on_right)
-        type = 'derived' if depends_on else 'source'
-        return ColumnLevelDependencyColumn(type=type, depends_on=depends_on)
-    elif isinstance(expression, Case):
-        ifs = expression.args['ifs']
-        default = expression.args['default']
-        depends_on = []
-        for expr in ifs:
-            depends_on_one = _cll_expression(expr, table_alias_map).depends_on
-            depends_on.extend(depends_on_one)
-        if default is not None:
-            depends_on.extend(_cll_expression(default, table_alias_map).depends_on)
-        type = 'derived' if depends_on else 'source'
-        return ColumnLevelDependencyColumn(type=type, depends_on=depends_on)
-    elif isinstance(expression, If):
-        depends_on = []
-        if expression.this:
-            depends_on_one = _cll_expression(expression.this, table_alias_map).depends_on
-            depends_on.extend(depends_on_one)
-        if expression.args.get('true'):
-            depends_on_one = _cll_expression(expression.args.get('true'), table_alias_map).depends_on
-            depends_on.extend(depends_on_one)
-        if expression.args.get('false'):
-            depends_on_one = _cll_expression(expression.args.get('false'), table_alias_map).depends_on
-            depends_on.extend(depends_on_one)
-        type = 'derived' if depends_on else 'source'
-        return ColumnLevelDependencyColumn(type=type, depends_on=depends_on)
-    elif isinstance(expression, Func):
-        if expression.expressions:
-            depends_on = []
-            for expr in expression.expressions:
-                depends_on_one = _cll_expression(expr, table_alias_map).depends_on
-                depends_on.extend(depends_on_one)
-            type = 'derived' if depends_on else 'source'
-            return ColumnLevelDependencyColumn(type=type, depends_on=depends_on)
-        if expression.this:
-            depends_on = _cll_expression(expression.this, table_alias_map).depends_on
-            type = 'derived' if depends_on else 'source'
-            return ColumnLevelDependencyColumn(type=type, depends_on=depends_on)
-
-        return ColumnLevelDependencyColumn(type='source', depends_on=[])
-    elif expression.this and isinstance(expression.this, Expression):
-        depends_on = _cll_expression(expression.this, table_alias_map).depends_on
-        type = 'derived' if depends_on else 'source'
-        return ColumnLevelDependencyColumn(type=type, depends_on=depends_on)
-    elif expression.expressions:
-        depends_on = []
-        for expr in expression.expressions:
-            depends_on_one = _cll_expression(expr, table_alias_map).depends_on
-            depends_on.extend(depends_on_one)
-        type = 'derived' if depends_on else 'source'
-        return ColumnLevelDependencyColumn(type=type, depends_on=depends_on)
-    else:
-        depends_on = []
-        return ColumnLevelDependencyColumn(type='source', depends_on=depends_on)
-
-
-def cll(sql, schema=None, dialect=None) -> Dict[str, ColumnLevelDependencyColumn]:
-    # given a sql, return the columns depends on
+
+    def selected_column_dependency(ref_column: exp.Column) -> Optional[CllColumn]:
+        column_name = ref_column.name
+        return c2c_map.get(column_name)
+
+    # joins clause: Reference the source columns
+    if select.args.get("joins"):
+        joins = select.args.get("joins")
+        for join in joins:
+            if isinstance(join, exp.Join):
+                for ref_column in join.find_all(exp.Column):
+                    if source_column_dependency(ref_column) is not None:
+                        m2c.extend(source_column_dependency(ref_column).depends_on)
+
+    # where clauses: Reference the source columns
+    if select.args.get("where"):
+        where = select.args.get("where")
+        if isinstance(where, exp.Where):
+            for ref_column in where.find_all(exp.Column):
+                if source_column_dependency(ref_column) is not None:
+                    m2c.extend(source_column_dependency(ref_column).depends_on)
+            for subquery in where.find_all(exp.Subquery):
+                sub_cll = subquery_cll(subquery)
+                if sub_cll is not None:
+                    sub_m2c, sub_c2c_map = sub_cll
+                    m2c.extend(sub_m2c)
+                    for sub_c in sub_c2c_map.values():
+                        m2c.extend(sub_c.depends_on)
+
+    # group by clause: Reference the source columns, column index
+    if select.args.get("group"):
+        group = select.args.get("group")
+        if isinstance(group, exp.Group):
+            for ref_column in group.find_all(exp.Column):
+                if source_column_dependency(ref_column) is not None:
+                    m2c.extend(source_column_dependency(ref_column).depends_on)
+
+    # having clause: Reference the source columns, selected columns
+    if select.args.get("having"):
+        having = select.args.get("having")
+        if isinstance(having, exp.Having):
+            for ref_column in having.find_all(exp.Column):
+                if source_column_dependency(ref_column) is not None:
+                    m2c.extend(source_column_dependency(ref_column).depends_on)
+                elif selected_column_dependency(ref_column) is not None:
+                    m2c.extend(selected_column_dependency(ref_column).depends_on)
+            for subquery in having.find_all(exp.Subquery):
+                sub_cll = subquery_cll(subquery)
+                if sub_cll is not None:
+                    sub_m2c, sub_c2c_map = sub_cll
+                    m2c.extend(sub_m2c)
+                    for sub_c in sub_c2c_map.values():
+                        m2c.extend(sub_c.depends_on)
+
+    # order by clause: Reference the source columns, selected columns, column index
+    if select.args.get("order"):
+        order = select.args.get("order")
+        if isinstance(order, exp.Order):
+            for ref_column in order.find_all(exp.Column):
+                if source_column_dependency(ref_column) is not None:
+                    m2c.extend(source_column_dependency(ref_column).depends_on)
+                elif selected_column_dependency(ref_column) is not None:
+                    m2c.extend(selected_column_dependency(ref_column).depends_on)
+
+    for source in scope.sources.values():
+        scope_cll_result = scope_cll_map.get(source)
+        if scope_cll_result is None:
+            continue
+        sub_m2c, _ = scope_cll_result
+        m2c.extend(sub_m2c)
+
+    m2c = _dedeup_depends_on(m2c)
+
+    return m2c, c2c_map
+
+
+def cll(sql, schema=None, dialect=None) -> CllResult:
+    # given a sql, return the cll for the sql
     # {
-    #   'column1': {
-    #      'transformation_type': 'transform' or 'original',
-    #      'depends_on': {
-    #           {node: model_id, column: column}, ...]
-    #      }
+    #     'depends_on': [{'node': 'model_id', 'column': 'column'}],
+    #     'columns': {
+    #         'column1': {
+    #             'type': 'derived',
+    #             'depends_on': [{'node': 'model_id', 'column': 'column'}],
+    #         }
+    #     }
+    # }
 
     dialect = Dialect.get(dialect) if dialect is not None else None
 
     try:
         expression = parse_one(sql, dialect=dialect)
     except SqlglotError as e:
-        raise RecceException(f'Failed to parse SQL: {str(e)}')
+        raise RecceException(f"Failed to parse SQL: {str(e)}")
 
     try:
         expression = qualify(expression, schema=schema, dialect=dialect)
     except OptimizeError as e:
-        raise RecceException(f'Failed to optimize SQL: {str(e)}')
+        raise RecceException(f"Failed to optimize SQL: {str(e)}")
     except SqlglotError as e:
-        raise RecceException(f'Failed to qualify SQL: {str(e)}')
+        raise RecceException(f"Failed to qualify SQL: {str(e)}")
 
-    result = {}
-    global_lineage = {}
+    result = None
+    scope_cll_map = {}
     for scope in traverse_scope(expression):
-        scope_lineage = {}
-
-        table_alias_map = {
-            t.alias_or_name: t.name
-            for t in scope.tables
-        }
-
-        if isinstance(scope.expression, Union) or isinstance(scope.expression, Intersect):
-            for union_scope in scope.union_scopes:
-                for k, v in global_lineage[union_scope].items():
-                    if k not in scope_lineage:
-                        scope_lineage[k] = v
-                    else:
-                        scope_lineage[k].depends_on.extend(v.depends_on)
-                        scope_lineage[k].type = 'derived'
+        scope_type = scope.expression.key
+        if scope_type == "union" or scope_type == "intersect" or scope_type == "except":
+            result = _cll_set_scope(scope, scope_cll_map)
+        elif scope_type == "select":
+            result = _cll_select_scope(scope, scope_cll_map)
         else:
-            for select in scope.expression.selects:
-                # instance of Column
-                if isinstance(select, Column):
-                    # 'select a'
-                    column = select
-                    column_cll = _cll_expression(column, table_alias_map)
-                elif isinstance(select, Alias):
-                    # 'select a as b'
-                    # 'select CURRENT_TIMESTAMP() as create_at'
-                    alias = select
-                    col_expression = alias.this
-                    column_cll = _cll_expression(col_expression, table_alias_map)
-                    if (
-                        column_cll and
-                        column_cll.type == 'passthrough' and
-                        column_cll.depends_on[0].column != alias.alias_or_name
-                    ):
-                        column_cll.type = 'renamed'
-                else:
-                    # 'select 1'
-                    column_cll = ColumnLevelDependencyColumn(type='source', depends_on=[])
-
-                cte_type = None
-                flatten_col_depends_on = []
-                for col_dep in column_cll.depends_on:
-                    col_dep_node = col_dep.node
-                    col_dep_column = col_dep.column
-                    # cte
-                    cte_scope = scope.cte_sources.get(col_dep_node)
-                    # inline derived table
-                    source_scope = None
-                    if isinstance(scope.sources.get(col_dep_node), Scope):
-                        source_scope = scope.sources.get(col_dep_node)
-
-                    if cte_scope is not None:
-                        cte_cll = global_lineage[cte_scope]
-                        if cte_cll is None or cte_cll.get(col_dep_column) is None:
-                            # In dbt-duckdb, the external source is compiled as `read_csv('..') rather than a table.
-                            continue
-                        cte_type = cte_cll.get(col_dep_column).type
-                        flatten_col_depends_on.extend(cte_cll.get(col_dep_column).depends_on)
-                    elif source_scope is not None:
-                        source_cll = global_lineage[source_scope]
-                        if source_cll is None or source_cll.get(col_dep_column) is None:
-                            continue
-                        flatten_col_depends_on.extend(source_cll.get(col_dep_column).depends_on)
-                    else:
-                        flatten_col_depends_on.append(col_dep)
-
-                # deduplicate
-                dedup_col_depends_on = []
-                dedup_set = set()
-                for col_dep in flatten_col_depends_on:
-                    node_col = col_dep.node + '.' + col_dep.column
-                    if node_col not in dedup_set:
-                        dedup_col_depends_on.append(col_dep)
-                        dedup_set.add(node_col)
-
-                # transformation type
-                type = column_cll.type
-                if type == 'derived':
-                    if len(dedup_col_depends_on) == 0:
-                        type = 'source'
-                    else:
-                        # keep current scope type
-                        pass
-                elif cte_type is not None:
-                    if len(dedup_col_depends_on) > 1:
-                        type = 'derived'
-                    elif len(dedup_col_depends_on) == 0:
-                        type = 'source'
-                    else:
-                        if isinstance(select, Column):
-                            type = cte_type
-                        elif isinstance(select, Alias):
-                            alias = select
-                            if column_cll.depends_on[0].column == alias.alias_or_name:
-                                type = cte_type
-                            else:
-                                type = 'renamed' if cte_type == 'passthrough' else cte_type
-                        else:
-                            type = 'source'
-
-                scope_lineage[select.alias_or_name] = ColumnLevelDependencyColumn(
-                    type=type,
-                    depends_on=dedup_col_depends_on
-                )
-
-        global_lineage[scope] = scope_lineage
-        if not scope.is_cte:
-            result = scope_lineage
+            continue
+
+        scope_cll_map[scope] = result
 
+    if result is None:
+        raise RecceException("Failed to extract CLL from SQL")
     return result

recce/util/io.py

@@ -1,14 +1,14 @@
 import gzip
 import os
 import tempfile
-from abc import ABCMeta, abstractmethod, ABC
+from abc import ABC, ABCMeta, abstractmethod
 from enum import Enum
 
 
 class SupportedFileTypes(Enum):
-    FILE = 'file'
-    GZIP = 'gzip'
-    ZIP = 'zip'
+    FILE = "file"
+    GZIP = "gzip"
+    ZIP = "zip"
 
 
 def file_io_factory(file_type: SupportedFileTypes):
@@ -19,7 +19,7 @@ def file_io_factory(file_type: SupportedFileTypes):
     elif file_type == SupportedFileTypes.ZIP:
         return ZipFileIO
     else:
-        raise ValueError(f'Unsupported file type: {file_type}')
+        raise ValueError(f"Unsupported file type: {file_type}")
 
 
 class AbstractFileIO(metaclass=ABCMeta):
@@ -37,24 +37,24 @@ class AbstractFileIO(metaclass=ABCMeta):
 class FileIO(AbstractFileIO, ABC):
     @staticmethod
     def write(path: str, data: str, **kwargs):
-        with open(path, 'w') as f:
+        with open(path, "w", encoding="utf-8") as f:
             f.write(data)
 
     @staticmethod
     def read(path: str, **kwargs) -> str:
-        with open(path, 'r') as f:
+        with open(path, "r", encoding="utf-8") as f:
             return f.read()
 
 
 class GzipFileIO(AbstractFileIO, ABC):
     @staticmethod
     def write(path: str, data: str, **kwargs):
-        with gzip.open(path, 'wt') as f:
+        with gzip.open(path, "wt") as f:
             f.write(data)
 
     @staticmethod
     def read(path: str, **kwargs) -> str:
-        with gzip.open(path, 'rt') as f:
+        with gzip.open(path, "rt") as f:
             return f.read()
 
     @staticmethod
@@ -68,18 +68,22 @@ class ZipFileIO(AbstractFileIO, ABC):
     def _is_pyminizip_installed():
         try:
             import pyminizip
+
+            # Use the module to avoid F401
+            return pyminizip is not None
         except ImportError:
-            raise ImportError('pyminizip is not installed. Please install it using `pip install pyminizip`')
+            raise ImportError("pyminizip is not installed. Please install it using `pip install pyminizip`")
 
     @staticmethod
     def read(path: str, **kwargs) -> str:
         ZipFileIO._is_pyminizip_installed()
         import pyminizip
+
         cwd = os.getcwd()
-        password = kwargs.get('password')
-        zip_dir_name = kwargs.get('zip_dir_name')
+        password = kwargs.get("password")
+        zip_dir_name = kwargs.get("zip_dir_name")
         if zip_dir_name is None:
-            raise ValueError('zip_dir_name is required for zipping')
+            raise ValueError("zip_dir_name is required for zipping")
 
         try:
             with tempfile.TemporaryDirectory() as tmp_dir:
@@ -88,7 +92,7 @@ class ZipFileIO(AbstractFileIO, ABC):
                 content = FileIO.read(tmp_file)
         except Exception as e:
             error_msg = str(e)
-            if '-3' in error_msg:
+            if "-3" in error_msg:
                 raise Exception("Invalid password to uncompress state file.")
             raise Exception(f"Failed to uncompress state file: {error_msg}")
         finally:
@@ -100,11 +104,12 @@ class ZipFileIO(AbstractFileIO, ABC):
     def write(path: str, data: str, **kwargs):
         ZipFileIO._is_pyminizip_installed()
         import pyminizip
+
         cwd = os.getcwd()
-        password = kwargs.get('password')
-        zip_dir_name = kwargs.get('zip_dir_name')
+        password = kwargs.get("password")
+        zip_dir_name = kwargs.get("zip_dir_name")
         if zip_dir_name is None:
-            raise ValueError('zip_dir_name is required for zipping')
+            raise ValueError("zip_dir_name is required for zipping")
 
         try:
             with tempfile.TemporaryDirectory() as tmp_dir: