diff --git a/py-polars/tests/unit/interop/test_interop.py b/py-polars/tests/unit/interop/test_interop.py index 8ae1b9a3b55e..daad87df371d 100644 --- a/py-polars/tests/unit/interop/test_interop.py +++ b/py-polars/tests/unit/interop/test_interop.py @@ -13,6 +13,7 @@ from polars.exceptions import ComputeError, UnstableWarning from polars.interchange.protocol import CompatLevel from polars.testing import assert_frame_equal, assert_series_equal +from tests.unit.utils.pycapsule_utils import PyCapsuleStreamHolder def test_arrow_list_roundtrip() -> None: @@ -752,26 +753,8 @@ def test_compat_level(monkeypatch: pytest.MonkeyPatch) -> None: def test_df_pycapsule_interface() -> None: - class PyCapsuleStreamHolder: - """ - Hold the Arrow C Stream pycapsule. - - A class that exposes _only_ the Arrow C Stream interface via Arrow PyCapsules. - This ensures that pyarrow is seeing _only_ the `__arrow_c_stream__` dunder, and - that nothing else (e.g. the dataframe or array interface) is actually being - used. - """ - - capsule: object - - def __init__(self, capsule: object) -> None: - self.capsule = capsule - - def __arrow_c_stream__(self, requested_schema: object) -> object: - return self.capsule - df = pl.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) - out = pa.table(PyCapsuleStreamHolder(df.__arrow_c_stream__(None))) + out = pa.table(PyCapsuleStreamHolder(df)) assert df.shape == out.shape assert df.schema.names() == out.schema.names diff --git a/py-polars/tests/unit/series/test_series.py b/py-polars/tests/unit/series/test_series.py index ebd0abc23752..71e217607168 100644 --- a/py-polars/tests/unit/series/test_series.py +++ b/py-polars/tests/unit/series/test_series.py @@ -29,6 +29,7 @@ ShapeError, ) from polars.testing import assert_frame_equal, assert_series_equal +from tests.unit.utils.pycapsule_utils import PyCapsuleStreamHolder if TYPE_CHECKING: from zoneinfo import ZoneInfo @@ -629,26 +630,8 @@ def test_arrow() -> None: def test_pycapsule_interface() -> None: - class PyCapsuleSeriesHolder: - """ - Hold the Arrow C Stream pycapsule. - - A class that exposes _only_ the Arrow C Stream interface via Arrow PyCapsules. - This ensures that pyarrow is seeing _only_ the `__arrow_c_stream__` dunder, and - that nothing else (e.g. the dataframe or array interface) is actually being - used. - """ - - capsule: object - - def __init__(self, capsule: object): - self.capsule = capsule - - def __arrow_c_stream__(self, requested_schema: object) -> object: - return self.capsule - a = pl.Series("a", [1, 2, 3, None]) - out = pa.chunked_array(PyCapsuleSeriesHolder(a.__arrow_c_stream__(None))) + out = pa.chunked_array(PyCapsuleStreamHolder(a)) out_arr = out.combine_chunks() assert out_arr == pa.array([1, 2, 3, None]) diff --git a/py-polars/tests/unit/utils/pycapsule_utils.py b/py-polars/tests/unit/utils/pycapsule_utils.py new file mode 100644 index 000000000000..3e7285fc0202 --- /dev/null +++ b/py-polars/tests/unit/utils/pycapsule_utils.py @@ -0,0 +1,22 @@ +from typing import Any + + +class PyCapsuleStreamHolder: + """ + Hold the Arrow C Stream pycapsule. + + A class that exposes _only_ the Arrow C Stream interface via Arrow PyCapsules. + This ensures that pyarrow is seeing _only_ the `__arrow_c_stream__` dunder, and + that nothing else (e.g. the dataframe or array interface) is actually being + used. + + This is used by tests across multiple files. + """ + + arrow_obj: Any + + def __init__(self, arrow_obj: object) -> None: + self.arrow_obj = arrow_obj + + def __arrow_c_stream__(self, requested_schema: object = None) -> object: + return self.arrow_obj.__arrow_c_stream__(requested_schema)