feat: Trim the number of columns we output to stdout (#294)

GoogleCloudPlatform · Aug 20, 2021 · 8567ba1 · 8567ba1
1 parent 51f14e0
commit 8567ba1
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 5 deletions.
diff --git a/data_validation/consts.py b/data_validation/consts.py
@@ -100,3 +100,18 @@
 
 # Ibis Object Info
 NUMERIC_DATA_TYPES = ["float64", "int32", "int64", "decimal"]
+
+FORMAT_TYPES = ["csv", "json", "table", "text"]
+
+# Text Result Handler column filter list
+COLUMN_FILTER_LIST = [
+    "aggregation_type",
+    "difference",
+    "end_time",
+    "labels",
+    "pct_threshold",
+    "run_id",
+    "source_agg_value",
+    "start_time",
+    "target_agg_value",
+]
diff --git a/data_validation/result_handlers/text.py b/data_validation/result_handlers/text.py
@@ -22,10 +22,13 @@
 Output validation report to text-based log
 """
 
+from data_validation import consts
+
 
 class TextResultHandler(object):
-    def __init__(self, format):
+    def __init__(self, format, cols_filter_list=consts.COLUMN_FILTER_LIST):
         self.format = format
+        self.cols_filter_list = cols_filter_list
 
     def print_formatted_(self, result_df):
         """
@@ -38,16 +41,22 @@ def print_formatted_(self, result_df):
             print(result_df.to_csv(index=False))
         elif self.format == "json":
             print(result_df.to_json(orient="index"))
-        elif self.format == "table":
-            print(result_df.to_markdown(tablefmt="fancy_grid"))
         else:
+            print(
+                result_df.drop(self.cols_filter_list, axis=1).to_markdown(
+                    tablefmt="fancy_grid"
+                )
+            )
+
+        if self.format not in consts.FORMAT_TYPES:
             error_msg = (
                 f"format [{self.format}] not supported, results printed in default(table) mode. "
                 f"Supported formats are [text, csv, json, table]"
             )
-            print(result_df.to_markdown(tablefmt="fancy_grid"))
             raise ValueError(error_msg)
 
+        return result_df
+
     def execute(self, config, result_df):
         self.print_formatted_(result_df)
         return result_df
diff --git a/tests/unit/result_handlers/test_text.py b/tests/unit/result_handlers/test_text.py
@@ -16,7 +16,6 @@
 
 from pandas import DataFrame
 
-
 SAMPLE_CONFIG = {}
 SAMPLE_RESULT_DATA = [
     {"table_name": "my_table", "count": 10},
@@ -53,3 +52,26 @@ def test_unsupported_result_format(module_under_test):
 
         handler_output = result_handler.execute(SAMPLE_CONFIG, result_df)
         assert handler_output["count"].sum() == result_df["count"].sum()
+
+
+def test_columns_to_print(module_under_test, capsys):
+    """Check for trimmed columns in grid print"""
+    table = [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]
+    result_df = DataFrame(table, columns=["A", "B", "C", "D"])
+    cols_filter_list = ["B", "D"]
+    format = "table-"
+    result_handler = module_under_test.TextResultHandler(format, cols_filter_list)
+    handler_output = result_handler.execute(SAMPLE_CONFIG, result_df)
+
+    grid_text = "││A│C││0│0│2││1│4│6││2│8│10│"
+    printed_text = capsys.readouterr().out
+    printed_text = (
+        printed_text.replace("\n", "")
+        .replace("'", "")
+        .replace(" ", "")
+        .replace("╒════╤═════╤═════╕", "")
+        .replace("╞════╪═════╪═════╡", "")
+        .replace("├────┼─────┼─────┤", "")
+        .replace("╘════╧═════╧═════╛", "")
+    )
+    assert printed_text == grid_text