Splitted io test into each format (#2727)

py-polars/src/dataframe.rs

@@ -288,7 +288,7 @@ impl PyDataFrame {
         row_oriented: bool,
         json_lines: bool,
     ) -> PyResult<()> {
-        let file = get_file_like(py_f, false)?;
+        let file = get_file_like(py_f, true)?;
 
         let r = match (pretty, row_oriented, json_lines) {
             (_, true, true) => panic!("{}", "only one of {row_oriented, json_lines} should be set"),

py-polars/tests/conftest.py

@@ -4,6 +4,8 @@
 
 import polars as pl
 
+IO_TEST_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "io"))
+
 EXAMPLES_DIR = os.path.abspath(
     os.path.join(
         os.path.dirname(__file__),
@@ -31,6 +33,16 @@
 )
 
 
+@pytest.fixture
+def io_test_dir() -> str:
+    return IO_TEST_DIR
+
+
+@pytest.fixture
+def examples_dir() -> str:
+    return EXAMPLES_DIR
+
+
 @pytest.fixture
 def foods_csv() -> str:
     return FOODS_CSV

py-polars/tests/io/.gitignore

@@ -0,0 +1,5 @@
+*.parquet
+*.csv
+*.avro
+*.ipc
+*.json

py-polars/tests/io/test_avro.py

@@ -0,0 +1,38 @@
+# flake8: noqa: W191,E101
+import os
+from typing import List
+
+import pytest
+
+import polars as pl
+from polars import io
+
+
+@pytest.fixture
+def example_df() -> pl.DataFrame:
+    return pl.DataFrame({"i64": [1, 2], "f64": [0.1, 0.2], "utf8": ["a", "b"]})
+
+
+@pytest.fixture
+def compressions() -> List[str]:
+    return ["uncompressed", "snappy", "deflate"]
+
+
+def test_from_to_buffer(example_df: pl.DataFrame, compressions: List[str]) -> None:
+    for compression in compressions:
+        buf = io.BytesIO()
+        example_df.to_avro(buf, compression=compression)  # type: ignore
+        buf.seek(0)
+        read_df = pl.read_avro(buf)
+        assert example_df.frame_equal(read_df)
+
+
+def test_from_to_file(
+    io_test_dir: str, example_df: pl.DataFrame, compressions: List[str]
+) -> None:
+    f = os.path.join(io_test_dir, "small.avro")
+
+    for compression in compressions:
+        example_df.to_avro(f, compression=compression)  # type: ignore
+        df_read = pl.read_avro(str(f))
+        assert example_df.frame_equal(df_read)