MAINT: Prepare patsy for pandas3 StringDtype

.pre-commit-config.yaml

@@ -7,8 +7,6 @@ repos:
     - id: trailing-whitespace
     - id: fix-byte-order-marker
     - id: destroyed-symlinks
-    - id: fix-encoding-pragma
-      args: ["--remove"]
     - id: mixed-line-ending
     - id: name-tests-test
       args: ["--pytest-test-first"]

patsy/missing.py

@@ -183,7 +183,7 @@ def _handle_NA_drop(self, values, is_NAs, origins):
             total_mask |= is_NA
         good_mask = ~total_mask
         # "..." to handle 1- versus 2-dim indexing
-        return [v[good_mask, ...] for v in values]
+        return [v[good_mask] if v.ndim == 1 else v[good_mask, ...] for v in values]
 
     __getstate__ = no_pickling
 

patsy/test_highlevel.py

@@ -979,3 +979,21 @@ def test_C_and_pandas_categorical():
         assert np.allclose(
             dmatrix("C(obj, levels=['a', 'b'])", d), [[1, 0], [1, 1], [1, 0]]
         )
+
+
+def test_NAActioon_pandas_string_index():
+    if not have_pandas:
+        return
+    from patsy.missing import NAAction
+
+    formula = "1 + x + z"
+    action = NAAction("drop")
+    data = pandas.DataFrame(
+        {"z": [1.0, np.nan, 2.0], "x": [1, 2, 3]}, index=["a", "b", "c"]
+    )
+    dm = dmatrix(formula, data, 0, NA_action=action, return_type="dataframe")
+    di = dm.design_info
+    data2 = pandas.DataFrame({"z": [4.0, 5.0], "x": [6, 7]})
+    dm2 = dmatrix(di, data2, 0, return_type="dataframe")
+    assert np.allclose(dm2, [[1.0, 6.0, 4.0], [1.0, 7.0, 5.0]])
+    assert list(dm2.columns) == ["Intercept", "x", "z"]

patsy/util.py

@@ -40,14 +40,19 @@
 try:
     import pandas
 except ImportError:
-    have_pandas = False
+    PANDAS3 = have_pandas = False
 else:
     have_pandas = True
+    import packaging.version
+
+    pandas_version = packaging.version.parse(pandas.__version__)
+    PANDAS3 = pandas_version >= packaging.version.parse("3.0.0.dev0")
 
 # Pandas versions < 0.9.0 don't have Categorical
 # Can drop this guard whenever we drop support for such older versions of
 # pandas.
 have_pandas_categorical = have_pandas and hasattr(pandas, "Categorical")
+have_pandas_string_dtype = have_pandas and hasattr(pandas, "StringDtype")
 if not have_pandas:
     _pandas_is_categorical_dtype = None
 else:
@@ -65,6 +70,11 @@
         )
 have_pandas_categorical_dtype = _pandas_is_categorical_dtype is not None
 
+
+def safe_is_pandas_string_dtype(x):
+    return have_pandas_string_dtype and isinstance(x, pandas.StringDtype)
+
+
 # The handling of the `copy` keyword has been changed since numpy>=2.
 # https://numpy.org/devdocs/numpy_2_0_migration_guide.html#adapting-to-changes-in-the-copy-keyword
 # If numpy<2 support is dropped, this try-clause can be removed.
@@ -118,7 +128,9 @@ def test_asarray_or_pandas():
         assert s_view1.name == "A"
         assert np.array_equal(s_view1.index, [10, 20, 30])
         s_view1[10] = 101
-        assert s[10] == 101
+        # pandas 3 uses copy-on-write, so no longer valid
+        if not PANDAS3:
+            assert s[10] == 101
         s_copy = asarray_or_pandas(s, copy=True)
         assert s_copy.name == "A"
         assert np.array_equal(s_copy.index, [10, 20, 30])
@@ -130,14 +142,18 @@ def test_asarray_or_pandas():
         assert s_view2.name == "A"
         assert np.array_equal(s_view2.index, [10, 20, 30])
         s_view2[10] = 99
-        assert s[10] == 99
+        # pandas 3 uses copy-on-write, so no longer valid
+        if not PANDAS3:
+            assert s[10] == 99
 
         df = pandas.DataFrame([[1, 2, 3]], columns=["A", "B", "C"], index=[10])
         df_view1 = asarray_or_pandas(df)
         df_view1.loc[10, "A"] = 101
         assert np.array_equal(df_view1.columns, ["A", "B", "C"])
         assert np.array_equal(df_view1.index, [10])
-        assert df.loc[10, "A"] == 101
+        # pandas 3 uses copy-on-write, so no longer valid
+        if not PANDAS3:
+            assert df.loc[10, "A"] == 101
         df_copy = asarray_or_pandas(df, copy=True)
         assert np.array_equal(df_copy, df)
         assert np.array_equal(df_copy.columns, ["A", "B", "C"])
@@ -799,7 +815,8 @@ def test_safe_is_pandas_categorical():
 #   https://github.com/pydata/pandas/issues/9581
 #   https://github.com/pydata/pandas/issues/9581#issuecomment-77099564
 def safe_issubdtype(dt1, dt2):
-    if safe_is_pandas_categorical_dtype(dt1):
+    # The second condition is needed to support pandas >= 3 (!)
+    if safe_is_pandas_categorical_dtype(dt1) or safe_is_pandas_string_dtype(dt1):
         return False
     return np.issubdtype(dt1, dt2)
 

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = {py36,py37,py38,py39,py310,py311,py312,py313}-{with_pandas,without_pandas}
+envlist = {py36,py37,py38,py39,py310,py311,py312,py313,py314}-{with_pandas,without_pandas}
 
 [gh-actions]
 python =
@@ -11,6 +11,7 @@ python =
   3.11: py311
   3.12: py312
   3.13: py313
+  3.14: py314
 
 [testenv]
 deps=