feat: update standardize_dataframe to accept custom standard columns and adjust related tests

Author: bencap

src/mavedb/lib/validation/dataframe/dataframe.py

@@ -82,8 +82,8 @@ def validate_and_standardize_dataframe_pair(
     if not targets:
         raise ValueError("Can't validate provided file with no targets.")
 
-    standardized_scores_df = standardize_dataframe(scores_df)
-    standardized_counts_df = standardize_dataframe(counts_df) if counts_df is not None else None
+    standardized_scores_df = standardize_dataframe(scores_df, STANDARD_COLUMNS)
+    standardized_counts_df = standardize_dataframe(counts_df, STANDARD_COLUMNS) if counts_df is not None else None
 
     validate_dataframe(standardized_scores_df, "scores", targets, hdp)
 
@@ -224,7 +224,7 @@ def standardize_dict_keys(d: dict[str, Any]) -> dict[str, Any]:
     return {clean_col_name(k): v for k, v in d.items()}
 
 
-def standardize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
+def standardize_dataframe(df: pd.DataFrame, standard_columns: tuple[str, ...]) -> pd.DataFrame:
     """Standardize a dataframe by sorting the columns and changing the standard column names to lowercase.
     Also strips leading and trailing whitespace from column names and removes any quoted strings from column names.
 
@@ -250,7 +250,7 @@ def standardize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
     cleaned_columns = {c: clean_col_name(c) for c in df.columns}
     df.rename(columns=cleaned_columns, inplace=True)
 
-    column_mapper = {x: x.lower() for x in df.columns if x.lower() in STANDARD_COLUMNS}
+    column_mapper = {x: x.lower() for x in df.columns if x.lower() in standard_columns}
     df.rename(columns=column_mapper, inplace=True)
 
     return sort_dataframe_columns(df)

tests/validation/dataframe/test_dataframe.py

@@ -13,6 +13,7 @@
     required_score_column,
 )
 from mavedb.lib.validation.dataframe.dataframe import (
+    STANDARD_COLUMNS,
     choose_dataframe_index_column,
     sort_dataframe_columns,
     standardize_dataframe,
@@ -93,32 +94,36 @@ def test_sort_dataframe_preserves_extras_order(self):
 
 class TestStandardizeDataframe(DfTestCase):
     def test_preserve_standardized(self):
-        standardized_df = standardize_dataframe(self.dataframe)
+        standardized_df = standardize_dataframe(self.dataframe, STANDARD_COLUMNS)
         pd.testing.assert_frame_equal(self.dataframe, standardized_df)
 
     def test_standardize_changes_case_variants(self):
-        standardized_df = standardize_dataframe(self.dataframe.rename(columns={hgvs_nt_column: hgvs_nt_column.upper()}))
+        standardized_df = standardize_dataframe(
+            self.dataframe.rename(columns={hgvs_nt_column: hgvs_nt_column.upper()}), STANDARD_COLUMNS
+        )
         pd.testing.assert_frame_equal(self.dataframe, standardized_df)
 
     def test_standardize_changes_case_scores(self):
         standardized_df = standardize_dataframe(
-            self.dataframe.rename(columns={required_score_column: required_score_column.title()})
+            self.dataframe.rename(columns={required_score_column: required_score_column.title()}), STANDARD_COLUMNS
         )
         pd.testing.assert_frame_equal(self.dataframe, standardized_df)
 
     def test_standardize_preserves_extras_case(self):
-        standardized_df = standardize_dataframe(self.dataframe.rename(columns={"extra": "extra".upper()}))
+        standardized_df = standardize_dataframe(
+            self.dataframe.rename(columns={"extra": "extra".upper()}), STANDARD_COLUMNS
+        )
         pd.testing.assert_frame_equal(self.dataframe.rename(columns={"extra": "extra".upper()}), standardized_df)
 
     def test_standardize_removes_quotes(self):
         standardized_df = standardize_dataframe(
-            self.dataframe.rename(columns={"extra": "'extra'", "extra2": '"extra2"'})
+            self.dataframe.rename(columns={"extra": "'extra'", "extra2": '"extra2"'}), STANDARD_COLUMNS
         )
         pd.testing.assert_frame_equal(self.dataframe, standardized_df)
 
     def test_standardize_removes_whitespace(self):
         standardized_df = standardize_dataframe(
-            self.dataframe.rename(columns={"extra": " extra ", "extra2": "    extra2"})
+            self.dataframe.rename(columns={"extra": " extra ", "extra2": "    extra2"}), STANDARD_COLUMNS
         )
         pd.testing.assert_frame_equal(self.dataframe, standardized_df)
 
@@ -135,7 +140,8 @@ def test_standardize_sorts_columns(self):
                     "count1",
                     "extra",
                 ],
-            ]
+            ],
+            STANDARD_COLUMNS,
         )
         pd.testing.assert_frame_equal(
             self.dataframe[