wip

Author: MrPowers

chispa/default_formats.py

@@ -4,5 +4,5 @@
 class DefaultFormats:
     mismatched_rows = ["red"]
     matched_rows = ["blue"]
-    mismatched_cells = ["white", "underline"]
+    mismatched_cells = ["red", "underline"]
     matched_cells = ["blue"]

chispa/rows_comparer.py

@@ -40,61 +40,38 @@ def assert_basic_rows_equality(rows1, rows2, underline_cells=False, formats=Defa
             raise chispa.DataFramesNotEqualError("\n" + t.get_string())
 
 
-def assert_generic_rows_equality(rows1, rows2, row_equality_fun, row_equality_fun_args, underline_cells=False):
+def assert_generic_rows_equality(rows1, rows2, row_equality_fun, row_equality_fun_args, underline_cells=False, formats=DefaultFormats()):
     df1_rows = rows1
     df2_rows = rows2
     zipped = list(six.moves.zip_longest(df1_rows, df2_rows))
     t = PrettyTable(["df1", "df2"])
-    allRowsEqual = True
-    if underline_cells:
-        row_column_names = rows1[0].__fields__
-        num_columns = len(row_column_names)
+    all_rows_equal = True
     for r1, r2 in zipped:
         # rows are not equal when one is None and the other isn't
         if (r1 is not None and r2 is None) or (r2 is not None and r1 is None):
-            allRowsEqual = False
-            t.add_row([r1, r2])
+            all_rows_equal = False
+            t.add_row([format_string(r1, formats.mismatched_rows), format_string(r2, formats.mismatched_rows)])
         # rows are equal
         elif row_equality_fun(r1, r2, *row_equality_fun_args):
-            first = bcolors.LightBlue + str(r1) + bcolors.LightRed
-            second = bcolors.LightBlue + str(r2) + bcolors.LightRed
-            t.add_row([first, second])
+            r1_string = ", ".join(map(lambda f: f"{f}={r1[f]}", r1.__fields__))
+            r2_string = ", ".join(map(lambda f: f"{f}={r2[f]}", r2.__fields__))
+            t.add_row([format_string(r1_string, formats.matched_rows), format_string(r2_string, formats.matched_rows)])
         # otherwise, rows aren't equal
         else:
-            allRowsEqual = False
-            # Underline cells if requested
-            if underline_cells:
-                t.add_row(__underline_cells_in_row(
-                    r1=r1, r2=r2, row_column_names=row_column_names, num_columns=num_columns))
-            else:
-                t.add_row([r1, r2])
-    if allRowsEqual == False:
-        raise chispa.DataFramesNotEqualError("\n" + t.get_string())
-
-
-def __underline_cells_in_row(r1=Row, r2=Row, row_column_names=List[str], num_columns=int) -> List[str]:
-    """
-    Takes two Row types, a list of column names for the Rows and the length of columns
-    Returns list of two strings, with underlined columns within rows that are different for PrettyTable
-    """
-    r1_string = "Row("
-    r2_string = "Row("
-    for index, column in enumerate(row_column_names):
-        if ((index+1) == num_columns):
-            append_str = ""
-        else:
-            append_str = ", "
+            r_zipped = list(six.moves.zip_longest(r1.__fields__, r2.__fields__))
+            r1_string = []
+            r2_string = []
+            for r1_field, r2_field in r_zipped:
+                if r1[r1_field] != r2[r2_field]:
+                    all_rows_equal = False
+                    r1_string.append(format_string(f"{r1_field}='{r1[r1_field]}'", formats.mismatched_cells))
+                    r2_string.append(format_string(f"{r2_field}='{r2[r2_field]}'", formats.mismatched_cells))
+                else:
+                    r1_string.append(format_string(f"{r1_field}='{r1[r1_field]}'", formats.matched_cells))
+                    r2_string.append(format_string(f"{r2_field}='{r2[r2_field]}'", formats.matched_cells))
+            r1_res = ", ".join(r1_string)
+            r2_res = ", ".join(r2_string)
 
-        if r1[column] != r2[column]:
-            r1_string += underline_text(
-                f"{column}='{r1[column]}'") + f"{append_str}"
-            r2_string += underline_text(
-                f"{column}='{r2[column]}'") + f"{append_str}"
-        else:
-            r1_string += f"{column}='{r1[column]}'{append_str}"
-            r2_string += f"{column}='{r2[column]}'{append_str}"
-
-    r1_string += ")"
-    r2_string += ")"
-
-    return [bcolors.LightRed + r1_string, r2_string]
+            t.add_row([r1_res, r2_res])
+    if all_rows_equal == False:
+        raise chispa.DataFramesNotEqualError("\n" + t.get_string())

tests/test_readme_examples.py

@@ -151,8 +151,8 @@ def it_prints_underline_message():
             ("rick", 66),
         ]
         df2 = spark.createDataFrame(data, ["firstname", "age"]) 
-        # with pytest.raises(DataFramesNotEqualError) as e_info:
-        #     assert_df_equality(df1, df2, formats=my_formats)
+        with pytest.raises(DataFramesNotEqualError) as e_info:
+            assert_df_equality(df1, df2, underline_cells=True)
 
     def it_shows_assert_basic_rows_equality(my_formats):
         data = [
@@ -231,6 +231,7 @@ def test_approx_df_equality_different():
             (None, None)
         ]
         df2 = spark.createDataFrame(data2, ["num", "letter"])
+        # assert_approx_df_equality(df1, df2, 0.1)
         with pytest.raises(DataFramesNotEqualError) as e_info:
             assert_approx_df_equality(df1, df2, 0.1)