Merge pull request scikit-learn-contrib#393 from pimlock/fix-basen-in…

…verse-transform Fix basen_to_integer when column name contains regex metachar
bollwyvl · Jan 13, 2023 · a745057 · a745057
2 parents 1def428 + 7a4482e
commit a745057
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/category_encoders/basen.py b/category_encoders/basen.py
@@ -266,7 +266,7 @@ def basen_to_integer(self, X, cols, base):
         out_cols = X.columns.values.tolist()
 
         for col in cols:
-            col_list = [col0 for col0 in out_cols if re.match(str(col)+'_\\d+', str(col0))]
+            col_list = [col0 for col0 in out_cols if re.match(re.escape(str(col))+'_\\d+', str(col0))]
             insert_at = out_cols.index(col_list[0])
 
             if base == 1:

diff --git a/tests/test_basen.py b/tests/test_basen.py
@@ -140,6 +140,16 @@ def test_inverse_transform_HaveHandleMissingValueAndHandleUnknownReturnNan_Expec
 
         pd.testing.assert_frame_equal(expected, original)
 
+    def test_inverse_transform_HaveRegexMetacharactersInColumnName_ExpectInversed(self):
+        train = pd.DataFrame({'state (2-letter code)': ['il', 'ny', 'ca']})
+
+        enc = encoders.BaseNEncoder()
+        enc.fit(train)
+        result = enc.transform(train)
+        original = enc.inverse_transform(result)
+
+        pd.testing.assert_frame_equal(train, original)
+
     def test_num_cols(self):
         """
         Test that BaseNEncoder produces the correct number of output columns.