xfail based on version

Author: mroeschke

pandas/tests/generic/test_to_xarray.py

@@ -10,8 +10,9 @@
     date_range,
 )
 import pandas._testing as tm
+from pandas.util.version import Version
 
-pytest.importorskip("xarray")
+xarray = pytest.importorskip("xarray")
 
 
 class TestDataFrameToXArray:
@@ -30,13 +31,17 @@ def df(self):
             }
         )
 
-    def test_to_xarray_index_types(self, index_flat, df, using_infer_string):
+    def test_to_xarray_index_types(self, index_flat, df, request):
         index = index_flat
         # MultiIndex is tested in test_to_xarray_with_multiindex
         if len(index) == 0:
             pytest.skip("Test doesn't make sense for empty index")
-
-        from xarray import Dataset
+        elif Version(xarray.__version__) <= Version("2024.9.0"):
+            request.applymarker(
+                pytest.mark.xfail(
+                    reason="Categorical column not preserved.",
+                )
+            )
 
         df.index = index[:4]
         df.index.name = "foo"
@@ -46,7 +51,7 @@ def test_to_xarray_index_types(self, index_flat, df, using_infer_string):
         assert len(result.coords) == 1
         assert len(result.data_vars) == 8
         tm.assert_almost_equal(list(result.coords.keys()), ["foo"])
-        assert isinstance(result, Dataset)
+        assert isinstance(result, xarray.Dataset)
 
         # idempotency
         # datetimes w/tz are preserved
@@ -56,16 +61,12 @@ def test_to_xarray_index_types(self, index_flat, df, using_infer_string):
         tm.assert_frame_equal(result.to_dataframe(), expected)
 
     def test_to_xarray_empty(self, df):
-        from xarray import Dataset
-
         df.index.name = "foo"
         result = df[0:0].to_xarray()
         assert result.sizes["foo"] == 0
-        assert isinstance(result, Dataset)
+        assert isinstance(result, xarray.Dataset)
 
     def test_to_xarray_with_multiindex(self, df, using_infer_string):
-        from xarray import Dataset
-
         # MultiIndex
         df.index = MultiIndex.from_product([["a"], range(4)], names=["one", "two"])
         result = df.to_xarray()
@@ -74,7 +75,7 @@ def test_to_xarray_with_multiindex(self, df, using_infer_string):
         assert len(result.coords) == 2
         assert len(result.data_vars) == 8
         tm.assert_almost_equal(list(result.coords.keys()), ["one", "two"])
-        assert isinstance(result, Dataset)
+        assert isinstance(result, xarray.Dataset)
 
         result = result.to_dataframe()
         expected = df.copy()
@@ -88,7 +89,11 @@ def test_to_xarray_with_multiindex(self, df, using_infer_string):
 class TestSeriesToXArray:
     def test_to_xarray_index_types(self, index_flat, request):
         index = index_flat
-        if isinstance(index.dtype, StringDtype) and index.dtype.storage == "pyarrow":
+        if (
+            isinstance(index.dtype, StringDtype)
+            and index.dtype.storage == "pyarrow"
+            and Version(xarray.__version__) > Version("2024.9.0")
+        ):
             request.applymarker(
                 pytest.mark.xfail(
                     reason="xarray calling reshape of ArrowExtensionArray",
@@ -97,39 +102,33 @@ def test_to_xarray_index_types(self, index_flat, request):
             )
         # MultiIndex is tested in test_to_xarray_with_multiindex
 
-        from xarray import DataArray
-
         ser = Series(range(len(index)), index=index, dtype="int64")
         ser.index.name = "foo"
         result = ser.to_xarray()
         repr(result)
         assert len(result) == len(index)
         assert len(result.coords) == 1
         tm.assert_almost_equal(list(result.coords.keys()), ["foo"])
-        assert isinstance(result, DataArray)
+        assert isinstance(result, xarray.DataArray)
 
         # idempotency
         tm.assert_series_equal(result.to_series(), ser)
 
     def test_to_xarray_empty(self):
-        from xarray import DataArray
-
         ser = Series([], dtype=object)
         ser.index.name = "foo"
         result = ser.to_xarray()
         assert len(result) == 0
         assert len(result.coords) == 1
         tm.assert_almost_equal(list(result.coords.keys()), ["foo"])
-        assert isinstance(result, DataArray)
+        assert isinstance(result, xarray.DataArray)
 
     def test_to_xarray_with_multiindex(self):
-        from xarray import DataArray
-
         mi = MultiIndex.from_product([["a", "b"], range(3)], names=["one", "two"])
         ser = Series(range(6), dtype="int64", index=mi)
         result = ser.to_xarray()
         assert len(result) == 2
         tm.assert_almost_equal(list(result.coords.keys()), ["one", "two"])
-        assert isinstance(result, DataArray)
+        assert isinstance(result, xarray.DataArray)
         res = result.to_series()
         tm.assert_series_equal(res, ser)