black -l 120

dlawin · dlawin · commit c0375ac03afc · 2023-06-21T17:31:20.000-06:00
diff --git a/data_diff/dbt.py b/data_diff/dbt.py
@@ -71,7 +71,6 @@ def dbt_diff(
     config = dbt_parser.get_datadiff_config()
     _initialize_events(dbt_parser.dbt_user_id, dbt_parser.dbt_version, dbt_parser.dbt_project_id)
 
-
     if not state and not (config.prod_database or config.prod_schema):
         doc_url = "https://docs.datafold.com/development_testing/open_source#configure-your-dbt-project"
         raise DataDiffDbtProjectVarsNotFoundError(
@@ -130,14 +129,17 @@ def dbt_diff(
                 _local_diff(diff_vars, json_output)
         else:
             if json_output:
-                print(json.dumps(
-                    jsonify_error(
-                        table1=diff_vars.prod_path,
-                        table2=diff_vars.dev_path,
-                        dbt_model=diff_vars.dbt_model,
-                        error="No primary key found. Add uniqueness tests, meta, or tags.",
-                    )
-                ), flush=True)
+                print(
+                    json.dumps(
+                        jsonify_error(
+                            table1=diff_vars.prod_path,
+                            table2=diff_vars.dev_path,
+                            dbt_model=diff_vars.dbt_model,
+                            error="No primary key found. Add uniqueness tests, meta, or tags.",
+                        )
+                    ),
+                    flush=True,
+                )
             else:
                 rich.print(
                     _diff_output_base(".".join(diff_vars.dev_path), ".".join(diff_vars.prod_path))
@@ -189,7 +191,6 @@ def _get_diff_vars(
     )
 
 
-
 def _get_prod_path_from_config(config, model, dev_database, dev_schema) -> Tuple[str, str]:
     # "custom" dbt config database
     if model.config.database:
@@ -233,12 +234,11 @@ def _local_diff(diff_vars: TDiffVars, json_output: bool = False) -> None:
     dev_qualified_str = ".".join(diff_vars.dev_path)
     prod_qualified_str = ".".join(diff_vars.prod_path)
     diff_output_str = _diff_output_base(dev_qualified_str, prod_qualified_str)
-    
+
     table1 = connect_to_table(
         diff_vars.connection, prod_qualified_str, tuple(diff_vars.primary_keys), diff_vars.threads
     )
     table2 = connect_to_table(diff_vars.connection, dev_qualified_str, tuple(diff_vars.primary_keys), diff_vars.threads)
-    
 
     table1_columns = table1.get_schema()
     try:
@@ -294,23 +294,29 @@ def _local_diff(diff_vars: TDiffVars, json_output: bool = False) -> None:
         try:
             list(diff)
         except Exception as e:
-            print(json.dumps(jsonify_error(
-                list(table1.table_path),
-                list(table2.table_path),
-                diff_vars.dbt_model,
-                str(e)
-            )), flush=True)
+            print(
+                json.dumps(
+                    jsonify_error(list(table1.table_path), list(table2.table_path), diff_vars.dbt_model, str(e))
+                ),
+                flush=True,
+            )
             return
 
-        print(json.dumps(
-            jsonify(
-                diff, 
-                dbt_model=diff_vars.dbt_model,
-                with_summary=True, with_columns={
-                    "added": columns_added,
-                    "removed": columns_removed,
-                    "changed": columns_type_changed,
-        })), flush=True)
+        print(
+            json.dumps(
+                jsonify(
+                    diff,
+                    dbt_model=diff_vars.dbt_model,
+                    with_summary=True,
+                    with_columns={
+                        "added": columns_added,
+                        "removed": columns_removed,
+                        "changed": columns_type_changed,
+                    },
+                )
+            ),
+            flush=True,
+        )
         return
 
     if list(diff):
@@ -466,7 +472,7 @@ def _initialize_events(dbt_user_id: Optional[str], dbt_version: Optional[str], d
 
 
 def _email_signup() -> None:
-    email_regex = r'^[\w\.\+-]+@[\w\.-]+\.\w+$'
+    email_regex = r"^[\w\.\+-]+@[\w\.-]+\.\w+$"
     prompt = "\nWould you like to be notified when a new data-diff version is available?\n\nEnter email or leave blank to opt out (we'll only ask once).\n"
 
     if bool_ask_for_email():
diff --git a/data_diff/format.py b/data_diff/format.py
@@ -5,7 +5,7 @@
 from data_diff.diff_tables import DiffResultWrapper
 
 
-def jsonify_error(table1: List[str], table2: List[str], dbt_model: str, error: str) -> 'FailedDiff':
+def jsonify_error(table1: List[str], table2: List[str], dbt_model: str, error: str) -> "FailedDiff":
     return FailedDiff(
         status="failed",
         model=dbt_model,
@@ -15,10 +15,12 @@ def jsonify_error(table1: List[str], table2: List[str], dbt_model: str, error: s
     ).json()
 
 
-def jsonify(diff: DiffResultWrapper,
-            dbt_model: str,
-            with_summary: bool = False,
-            with_columns: Optional[Dict[str, List[str]]] = None) -> 'JsonDiff':
+def jsonify(
+    diff: DiffResultWrapper,
+    dbt_model: str,
+    with_summary: bool = False,
+    with_columns: Optional[Dict[str, List[str]]] = None,
+) -> "JsonDiff":
     """
     Converts the diff result into a JSON-serializable format.
     Optionally add stats summary and schema diff.
@@ -35,7 +37,6 @@ def jsonify(diff: DiffResultWrapper,
 
     t1_exclusive_rows, t2_exclusive_rows, diff_rows = _group_rows(diff_info, schema)
 
-
     diff_rows_jsonified = []
     for row in diff_rows:
         diff_rows_jsonified.append(_jsonify_diff(row, key_columns))
@@ -47,11 +48,11 @@ def jsonify(diff: DiffResultWrapper,
     t2_exclusive_rows_jsonified = []
     for row in t2_exclusive_rows:
         t2_exclusive_rows_jsonified.append(_jsonify_exclusive(row, key_columns))
-    
+
     summary = None
     if with_summary:
         summary = _jsonify_diff_summary(diff.get_stats_dict())
-    
+
     columns = None
     if with_columns:
         columns = _jsonify_columns_diff(with_columns, list(key_columns))
@@ -60,11 +61,8 @@ def jsonify(diff: DiffResultWrapper,
         t1_exclusive_rows
         or t2_exclusive_rows
         or diff_rows
-        or with_columns and (
-            with_columns['added']
-            or with_columns['removed']
-            or with_columns['changed']
-        )
+        or with_columns
+        and (with_columns["added"] or with_columns["removed"] or with_columns["changed"])
     )
     return JsonDiff(
         status="success",
@@ -73,23 +71,20 @@ def jsonify(diff: DiffResultWrapper,
         dataset1=list(table1.table_path),
         dataset2=list(table2.table_path),
         rows=RowsDiff(
-            exclusive=ExclusiveDiff(
-                dataset1=t1_exclusive_rows_jsonified,
-                dataset2=t2_exclusive_rows_jsonified
-            ),
+            exclusive=ExclusiveDiff(dataset1=t1_exclusive_rows_jsonified, dataset2=t2_exclusive_rows_jsonified),
             diff=diff_rows_jsonified,
         ),
         summary=summary,
         columns=columns,
     ).json()
 
 
-
 @dataclass
 class JsonExclusiveRowValue:
     """
     Value of a single column in a row
     """
+
     isPK: bool
     value: Any
 
@@ -99,6 +94,7 @@ class JsonDiffRowValue:
     """
     Pair of diffed values for 2 rows with equal PKs
     """
+
     dataset1: Any
     dataset2: Any
     isDiff: bool
@@ -163,38 +159,40 @@ class RowsDiff:
 
 @dataclass
 class FailedDiff:
-    status: str # Literal ["failed"]
+    status: str  # Literal ["failed"]
     model: str
     dataset1: List[str]
     dataset2: List[str]
     error: str
 
-    version: str = '1.0.0'
+    version: str = "1.0.0"
+
 
 @dataclass
 class JsonDiff:
-    status: str # Literal ["success"]
-    result: str # Literal ["different", "identical"]
+    status: str  # Literal ["success"]
+    result: str  # Literal ["different", "identical"]
     model: str
     dataset1: List[str]
     dataset2: List[str]
     rows: RowsDiff
     summary: Optional[JsonDiffSummary]
     columns: Optional[JsonColumnsSummary]
 
-    version: str = '1.0.0'
+    version: str = "1.0.0"
 
 
-def _group_rows(diff_info: DiffResultWrapper, 
-                schema: List[str]) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[Dict[str, Any]]]:
+def _group_rows(
+    diff_info: DiffResultWrapper, schema: List[str]
+) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], List[Dict[str, Any]]]:
     t1_exclusive_rows = []
     t2_exclusive_rows = []
     diff_rows = []
 
     for row in diff_info.diff:
         row_w_schema = dict(zip(schema, row))
-        is_t1_exclusive = row_w_schema['is_exclusive_a']
-        is_t2_exclusive = row_w_schema['is_exclusive_b']
+        is_t1_exclusive = row_w_schema["is_exclusive_a"]
+        is_t2_exclusive = row_w_schema["is_exclusive_b"]
 
         if is_t1_exclusive:
             t1_exclusive_rows.append(row_w_schema)
@@ -204,83 +202,72 @@ def _group_rows(diff_info: DiffResultWrapper,
 
         else:
             diff_rows.append(row_w_schema)
-    
+
     return t1_exclusive_rows, t2_exclusive_rows, diff_rows
 
 
 def _jsonify_diff(row: Dict[str, Any], key_columns: List[str]) -> Dict[str, JsonDiffRowValue]:
     columns = collections.defaultdict(dict)
     for field, value in row.items():
-        if field in ('is_exclusive_a', 'is_exclusive_b'):
+        if field in ("is_exclusive_a", "is_exclusive_b"):
             continue
 
-        if field.startswith('is_diff_'):
-            column_name = field.replace('is_diff_', '')
-            columns[column_name]['isDiff'] = bool(value)
+        if field.startswith("is_diff_"):
+            column_name = field.replace("is_diff_", "")
+            columns[column_name]["isDiff"] = bool(value)
 
-        elif field.endswith('_a'):
-            column_name = field.replace('_a', '')
-            columns[column_name]['dataset1'] = value
-            columns[column_name]['isPK'] = column_name in key_columns
+        elif field.endswith("_a"):
+            column_name = field.replace("_a", "")
+            columns[column_name]["dataset1"] = value
+            columns[column_name]["isPK"] = column_name in key_columns
 
-        elif field.endswith('_b'):
-            column_name = field.replace('_b', '')
-            columns[column_name]['dataset2'] = value
-            columns[column_name]['isPK'] = column_name in key_columns
-    
-    return {
-        column: JsonDiffRowValue(**data)
-        for column, data in columns.items()
-    }
+        elif field.endswith("_b"):
+            column_name = field.replace("_b", "")
+            columns[column_name]["dataset2"] = value
+            columns[column_name]["isPK"] = column_name in key_columns
+
+    return {column: JsonDiffRowValue(**data) for column, data in columns.items()}
 
 
 def _jsonify_exclusive(row: Dict[str, Any], key_columns: List[str]) -> Dict[str, JsonExclusiveRowValue]:
     columns = collections.defaultdict(dict)
     for field, value in row.items():
-        if field in ('is_exclusive_a', 'is_exclusive_b'):
+        if field in ("is_exclusive_a", "is_exclusive_b"):
             continue
-        if field.startswith('is_diff_'):
+        if field.startswith("is_diff_"):
             continue
-        if field.endswith('_b') and row['is_exclusive_b']:
-            column_name = field.replace('_b', '')
-            columns[column_name]['isPK'] = column_name in key_columns
-            columns[column_name]['value'] = value
-        elif field.endswith('_a') and row['is_exclusive_a']:
-            column_name = field.replace('_a', '')
-            columns[column_name]['isPK'] = column_name in key_columns
-            columns[column_name]['value'] = value
-    return {
-        column: JsonExclusiveRowValue(**data)
-        for column, data in columns.items()
-    }
+        if field.endswith("_b") and row["is_exclusive_b"]:
+            column_name = field.replace("_b", "")
+            columns[column_name]["isPK"] = column_name in key_columns
+            columns[column_name]["value"] = value
+        elif field.endswith("_a") and row["is_exclusive_a"]:
+            column_name = field.replace("_a", "")
+            columns[column_name]["isPK"] = column_name in key_columns
+            columns[column_name]["value"] = value
+    return {column: JsonExclusiveRowValue(**data) for column, data in columns.items()}
 
 
 def _jsonify_diff_summary(stats_dict: dict) -> JsonDiffSummary:
     return JsonDiffSummary(
         rows=Rows(
-            total=Total(
-                dataset1=stats_dict["rows_A"],
-                dataset2=stats_dict["rows_B"]
-            ),
+            total=Total(dataset1=stats_dict["rows_A"], dataset2=stats_dict["rows_B"]),
             exclusive=ExclusiveRows(
                 dataset1=stats_dict["exclusive_A"],
                 dataset2=stats_dict["exclusive_B"],
             ),
             updated=stats_dict["updated"],
-            unchanged=stats_dict["unchanged"]
+            unchanged=stats_dict["unchanged"],
         ),
-        stats=Stats(
-            diffCounts=stats_dict["stats"]['diff_counts']
-        )
+        stats=Stats(diffCounts=stats_dict["stats"]["diff_counts"]),
     )
 
 
 def _jsonify_columns_diff(columns_diff: Dict[str, List[str]], key_columns: List[str]) -> JsonColumnsSummary:
     return JsonColumnsSummary(
         primaryKey=key_columns,
-        exclusive= ExclusiveColumns(
-            dataset2= list(columns_diff.get('added', [])),
-            dataset1= list(columns_diff.get('removed', [])),
+        exclusive=ExclusiveColumns(
+            dataset2=list(columns_diff.get("added", [])),
+            dataset1=list(columns_diff.get("removed", [])),
         ),
-        typeChanged=list(columns_diff.get('changed', [])),
-    )
+        typeChanged=list(columns_diff.get("changed", [])),
+    )
diff --git a/data_diff/info_tree.py b/data_diff/info_tree.py
@@ -30,8 +30,7 @@ def update_from_children(self, child_infos):
         # self.diff = list(chain(*[c.diff for c in child_infos]))
         self.diff_count = sum(c.diff_count for c in child_infos if c.diff_count is not None)
         self.is_diff = any(c.is_diff for c in child_infos)
-        self.diff_schema = next((child.diff_schema for child in child_infos if child.diff_schema is not None),
-                                None)
+        self.diff_schema = next((child.diff_schema for child in child_infos if child.diff_schema is not None), None)
         self.diff = sum((c.diff for c in child_infos if c.diff is not None), [])
 
         self.rowcounts = {
diff --git a/tests/test_format.py b/tests/test_format.py
