Merge pull request #206 from mystic74/fix/dataframe-mutation-categorical-nan

Fix DataFrame mutation bug by copying data on initialization

Author: tompollard

tableone/tableone.py

@@ -304,7 +304,7 @@ def _repr_html_(self) -> str:
     def initialize_core_attributes(self, data, columns, categorical, continuous, groupby,
                                    nonnormal, min_max, pval, pval_adjust, htest_name,
                                    htest, missing, ddof, rename, sort, limit, order,
-                                   label_suffix, decimals, smd, overall, row_percent, 
+                                   label_suffix, decimals, smd, overall, row_percent,
                                    dip_test, normal_test, tukey_test, pval_threshold,
                                    include_null, pval_digits, ttest_equal_var,
                                    show_histograms, clip_histograms):
@@ -348,6 +348,8 @@ def initialize_core_attributes(self, data, columns, categorical, continuous, gro
         self._warnings = {}
 
         if self._categorical and self._include_null:
+            # Create a copy to avoid mutating the user's original DataFrame
+            data = data.copy()
             data[self._categorical] = handle_categorical_nulls(data[self._categorical], self._categorical)
 
         self._groupbylvls = get_groups(data, self._groupby, self._order, self._reserved_columns)

tests/unit/test_tableone.py

@@ -1396,6 +1396,94 @@ def test_handle_categorical_nulls_does_not_affect_continuous():
     assert result['cat'].iloc[2] == 'None'
 
 
+def test_handle_categorical_nulls_does_not_mutate_original():
+    """
+    Test that handle_categorical_nulls does not mutate the original DataFrame.
+    This tests the preprocessor function directly.
+    """
+    df = pd.DataFrame({
+        'cat': [1, 2, np.nan],
+        'other': ['a', 'b', 'c']
+    })
+
+    df_original = df.copy()
+    result = handle_categorical_nulls(df, categorical=['cat'])
+
+    pd.testing.assert_frame_equal(df, df_original, check_dtype=True)
+    assert result['cat'].iloc[2] == 'None'
+
+
+def test_tableone_does_not_mutate_input_with_numeric_categorical_nan():
+    """
+    Test that TableOne does not mutate the original DataFrame when processing
+    numeric categorical variables with NaN values.
+
+    Regression test for issue where NaN values were converted to string 'None'
+    in the original DataFrame.
+    """
+    df = pd.DataFrame({
+        'id': [1, 2, 3],
+        'category_col': [1, np.nan, 2]
+    })
+
+    df_original = df.copy()
+
+    TableOne(df, categorical=['category_col'], pval=False)
+
+    pd.testing.assert_frame_equal(df, df_original, check_dtype=True)
+
+
+def test_tableone_does_not_mutate_input_with_string_categorical_nan():
+    """
+    Test that TableOne does not mutate the original DataFrame when processing
+    string categorical variables with NaN values.
+    """
+    df = pd.DataFrame({
+        'id': [1, 2, 3],
+        'category_col': ['A', np.nan, 'B']
+    })
+
+    df_original = df.copy()
+
+    TableOne(df, categorical=['category_col'], pval=False)
+
+    pd.testing.assert_frame_equal(df, df_original, check_dtype=True)
+
+
+def test_tableone_does_not_mutate_input_with_groupby_and_pval():
+    """
+    Test that TableOne does not mutate the original DataFrame when using
+    groupby and pval with categorical variables containing NaN.
+    """
+    df = pd.DataFrame({
+        'group': ['A', 'B', 'A', 'B'],
+        'cat_var': [1, 2, np.nan, 3]
+    })
+
+    df_original = df.copy()
+
+    TableOne(df, categorical=['cat_var'], groupby='group', pval=True)
+
+    pd.testing.assert_frame_equal(df, df_original, check_dtype=True)
+
+
+def test_tableone_does_not_mutate_input_with_include_null_false():
+    """
+    Test that TableOne does not mutate the original DataFrame even when
+    include_null=False.
+    """
+    df = pd.DataFrame({
+        'id': [1, 2, 3],
+        'category_col': [1, np.nan, 2]
+    })
+
+    df_original = df.copy()
+
+    TableOne(df, categorical=['category_col'], include_null=False, pval=False)
+
+    pd.testing.assert_frame_equal(df, df_original, check_dtype=True)
+
+
 def test_pval_digits_custom_formatting():
     df = pd.DataFrame({
         'group': ['A', 'A', 'B', 'B', 'B'],