Fix formatting and tests for 3.10 (hopefully)

Author: cameronwaterman

nisystemlink/clients/dataframe/_data_frame_client.py

@@ -3,7 +3,7 @@
 import pyarrow as pa  # type: ignore
 from collections.abc import Iterable
 from io import BytesIO
-from typing import List, Optional, Union
+from typing import Any, List, Optional, Union
 
 from nisystemlink.clients import core
 from nisystemlink.clients.core._uplink._base_client import BaseClient
@@ -252,24 +252,19 @@ def get_table_data(
         """
         ...
 
-    @post(
-        "tables/{id}/data",
-        args=[Path, Body]
-    )
+    @post("tables/{id}/data", args=[Path, Body])
     def _append_table_data_json(
         self, id: str, data: models.AppendTableDataRequest
-    ) -> None:
-        ...
+    ) -> None: ...
 
     @post(
         "tables/{id}/data",
         args=[Path, Body, Query("endOfData")],
         content_type="application/vnd.apache.arrow.stream",
     )
     def _append_table_data_arrow(
-        self, id: str, data: Iterable[bytes], end_of_data: Optional[bool] = None
-    ) -> None:
-        ...
+        self, id: str, data: bytes, end_of_data: Optional[bool] = None
+    ) -> None: ...
 
     def append_table_data(
         self,
@@ -278,7 +273,7 @@ def append_table_data(
             Union[
                 models.AppendTableDataRequest,
                 models.DataFrame,
-                Iterable["pa.RecordBatch"], # type: ignore[name-defined]
+                Iterable["pa.RecordBatch"],  # type: ignore[name-defined]
             ]
         ],
         *,
@@ -343,44 +338,34 @@ def append_table_data(
                     "Iterable provided to data must yield pyarrow.RecordBatch objects."
                 )
 
-            def _generate_body() -> Iterable[memoryview]:
-                data_iter = iter(data)
-                try:
-                    batch = next(data_iter)
-                except StopIteration:
-                    return
+            def _build_body() -> bytes:
                 with BytesIO() as buf:
                     options = pa.ipc.IpcWriteOptions(compression="zstd")
-                    writer = pa.ipc.new_stream(buf, batch.schema, options=options)
-
-                    while True:
-                        writer.write_batch(batch)
-                        with buf.getbuffer() as view, view[0 : buf.tell()] as slice:
-                            yield slice
-                        buf.seek(0)
-                        try:
-                            batch = next(data_iter)
-                        except StopIteration:
-                            break
-
-                    writer.close()
-                    with buf.getbuffer() as view:
-                        with view[0 : buf.tell()] as slice:
-                            yield slice
+                    with pa.ipc.new_stream(
+                        buf, first_batch.schema, options=options
+                    ) as writer:
+                        writer.write_batch(first_batch)
+                        for batch in iterator:
+                            writer.write_batch(batch)
+                    return buf.getvalue()
 
             try:
                 self._append_table_data_arrow(
                     id,
-                    _generate_body(),
+                    _build_body(),
                     (end_of_data if end_of_data is not None else None),
                 )
             except core.ApiException as ex:
                 if ex.http_status_code == 400:
                     wrap = True
                     try:
-                        write_op = getattr(self.api_info().operations, "write_data", None)
-                        if write_op is not None and getattr(write_op, "version", 0) >= 2:
-                            # Service claims Arrow-capable write version; re-raise original exception
+                        write_op = getattr(
+                            self.api_info().operations, "write_data", None
+                        )
+                        if (
+                            write_op is not None
+                            and getattr(write_op, "version", 0) >= 2
+                        ):
                             wrap = False
                     except Exception:
                         pass

tests/integration/dataframe/test_dataframe.py

@@ -89,8 +89,17 @@ def test_tables(create_table):
 class TestDataFrame:
     def _new_single_int_table(self, create_table, column_name: str = "a") -> str:
         return create_table(
-            CreateTableRequest(columns=[Column(name=column_name, data_type=DataType.Int64, column_type=ColumnType.Index)])
+            CreateTableRequest(
+                columns=[
+                    Column(
+                        name=column_name,
+                        data_type=DataType.Int64,
+                        column_type=ColumnType.Index,
+                    )
+                ]
+            )
         )
+
     def test__api_info__returns(self, client):
         response = client.api_info()
 
@@ -620,78 +629,123 @@ def test__export_table_data__works(self, client: DataFrameClient, create_table):
                 == b'"col1","col2","col3"\r\n1,2.5,6.5\r\n2,1.5,5.5\r\n3,2.5,7.5'
             )
 
-    def test__append_table_data__append_request_success(self, client: DataFrameClient, create_table):
+    def test__append_table_data__append_request_success(
+        self, client: DataFrameClient, create_table
+    ):
         table_id = self._new_single_int_table(create_table)
         frame = DataFrame(columns=["a"], data=[["1"], ["2"]])
-        client.append_table_data(table_id, AppendTableDataRequest(frame=frame, end_of_data=True))
+        client.append_table_data(
+            table_id, AppendTableDataRequest(frame=frame, end_of_data=True)
+        )
 
-    def test__append_table_data__append_request_with_end_of_data_argument_disallowed(self, client: DataFrameClient, create_table):
+    def test__append_table_data__append_request_with_end_of_data_argument_disallowed(
+        self, client: DataFrameClient, create_table
+    ):
         request = AppendTableDataRequest(end_of_data=True)
-        with pytest.raises(ValueError, match="end_of_data must not be provided separately when passing an AppendTableDataRequest."):
-            client.append_table_data(self._new_single_int_table(create_table), request, end_of_data=True)
+        with pytest.raises(
+            ValueError,
+            match="end_of_data must not be provided separately when passing an AppendTableDataRequest.",
+        ):
+            client.append_table_data(
+                self._new_single_int_table(create_table), request, end_of_data=True
+            )
 
-    def test__append_table_data__append_request_without_end_of_data_success(self, client: DataFrameClient, create_table):
+    def test__append_table_data__append_request_without_end_of_data_success(
+        self, client: DataFrameClient, create_table
+    ):
         table_id = self._new_single_int_table(create_table)
         frame = DataFrame(columns=["a"], data=[["7"], ["8"]])
         client.append_table_data(table_id, AppendTableDataRequest(frame=frame))
 
-    def test__append_table_data__accepts_dataframe_model(self, client: DataFrameClient, create_table):
+    def test__append_table_data__accepts_dataframe_model(
+        self, client: DataFrameClient, create_table
+    ):
         table_id = self._new_single_int_table(create_table)
         frame = DataFrame(columns=["a"], data=[["1"], ["2"]])
         client.append_table_data(table_id, frame, end_of_data=True)
 
-    def test__append_table_data__dataframe_without_end_of_data_success(self, client: DataFrameClient, create_table):
+    def test__append_table_data__dataframe_without_end_of_data_success(
+        self, client: DataFrameClient, create_table
+    ):
         table_id = self._new_single_int_table(create_table)
         frame = DataFrame(columns=["a"], data=[["10"], ["11"]])
         client.append_table_data(table_id, frame)
 
-    def test__append_table_data__none_without_end_of_data_raises(self, client: DataFrameClient, create_table):
+    def test__append_table_data__none_without_end_of_data_raises(
+        self, client: DataFrameClient, create_table
+    ):
         table_id = create_table(basic_table_model)
-        with pytest.raises(ValueError, match="end_of_data must be provided when data is None"):
+        with pytest.raises(
+            ValueError, match="end_of_data must be provided when data is None"
+        ):
             client.append_table_data(table_id, None)
 
-    def test__append_table_data__flush_only_with_none(self, client: DataFrameClient, create_table):
+    def test__append_table_data__flush_only_with_none(
+        self, client: DataFrameClient, create_table
+    ):
         table_id = self._new_single_int_table(create_table)
         frame = DataFrame(columns=["a"], data=[["1"]])
         client.append_table_data(table_id, frame)
         client.append_table_data(table_id, None, end_of_data=True)
 
-    def test__append_table_data__arrow_ingestion_success(self, client: DataFrameClient, create_table):
+    def test__append_table_data__arrow_ingestion_success(
+        self, client: DataFrameClient, create_table
+    ):
         pa = pytest.importorskip("pyarrow")
         table_id = self._new_single_int_table(create_table)
         batch = pa.record_batch([pa.array([10, 11, 12])], names=["a"])
         client.append_table_data(table_id, [batch], end_of_data=True)
         with pytest.raises(ApiException):
             client.append_table_data(table_id, None, end_of_data=True)
 
-    def test__append_table_data__arrow_ingestion_with_end_of_data_query_param_false(self, client: DataFrameClient, create_table):
+    def test__append_table_data__arrow_ingestion_with_end_of_data_query_param_false(
+        self, client: DataFrameClient, create_table
+    ):
         pa = pytest.importorskip("pyarrow")
         table_id = self._new_single_int_table(create_table)
         batch1 = pa.record_batch([pa.array([4, 5, 6])], names=["a"])
         client.append_table_data(table_id, [batch1], end_of_data=False)
         batch2 = pa.record_batch([pa.array([7, 8])], names=["a"])
         client.append_table_data(table_id, [batch2], end_of_data=True)
 
-    def test__append_table_data__empty_iterator_requires_end_of_data(self, client: DataFrameClient, create_table):
+    def test__append_table_data__empty_iterator_requires_end_of_data(
+        self, client: DataFrameClient, create_table
+    ):
         table_id = create_table(basic_table_model)
-        with pytest.raises(ValueError, match="end_of_data must be provided when data iterator is empty."):
+        with pytest.raises(
+            ValueError,
+            match="end_of_data must be provided when data iterator is empty.",
+        ):
             client.append_table_data(table_id, [])
         client.append_table_data(table_id, [], end_of_data=True)
 
-    def test__append_table_data__arrow_iterable_with_non_recordbatch_elements_raises(self, client: DataFrameClient, create_table):
+    def test__append_table_data__arrow_iterable_with_non_recordbatch_elements_raises(
+        self, client: DataFrameClient, create_table
+    ):
         pytest.importorskip("pyarrow")
         table_id = create_table(basic_table_model)
-        with pytest.raises(ValueError, match="Iterable provided to data must yield pyarrow.RecordBatch objects."):
+        with pytest.raises(
+            ValueError,
+            match="Iterable provided to data must yield pyarrow.RecordBatch objects.",
+        ):
             client.append_table_data(table_id, [1, 2, 3])
 
-    def test__append_table_data__arrow_iterable_without_pyarrow_raises_runtime_error(self, client: DataFrameClient, create_table, monkeypatch):
+    def test__append_table_data__arrow_iterable_without_pyarrow_raises_runtime_error(
+        self, client: DataFrameClient, create_table, monkeypatch
+    ):
         import nisystemlink.clients.dataframe._data_frame_client as df_module
+
         monkeypatch.setattr(df_module, "pa", None)
         table_id = create_table(basic_table_model)
-        with pytest.raises(RuntimeError, match="pyarrow is not installed. Install to stream RecordBatches."):
+        with pytest.raises(
+            RuntimeError,
+            match="pyarrow is not installed. Install to stream RecordBatches.",
+        ):
             client.append_table_data(table_id, [object()])
 
-    def test__append_table_data__arrow_ingestion_400_unsupported(self, client: DataFrameClient):
+    def test__append_table_data__arrow_ingestion_400_unsupported(
+        self, client: DataFrameClient
+    ):
         pa = pytest.importorskip("pyarrow")
         table_id = "mock_table_id"
         bad_batch = pa.record_batch([pa.array([1, 2, 3])], names=["b"])
@@ -722,7 +776,9 @@ def test__append_table_data__arrow_ingestion_400_unsupported(self, client: DataF
 
         assert "Arrow ingestion request was rejected" in str(excinfo.value)
 
-    def test__append_table_data__arrow_ingestion_400_supported_passthrough(self, client: DataFrameClient):
+    def test__append_table_data__arrow_ingestion_400_supported_passthrough(
+        self, client: DataFrameClient
+    ):
         pa = pytest.importorskip("pyarrow")
         table_id = "mock_table_id"
         bad_batch = pa.record_batch([pa.array([1, 2, 3])], names=["b"])
@@ -753,21 +809,25 @@ def test__append_table_data__arrow_ingestion_400_supported_passthrough(self, cli
 
         assert "Arrow ingestion request was rejected" not in str(excinfo.value)
 
-    def test__append_table_data__arrow_ingestion_non_400_passthrough(self, client: DataFrameClient):
+    def test__append_table_data__arrow_ingestion_non_400_passthrough(
+        self, client: DataFrameClient
+    ):
         pa = pytest.importorskip("pyarrow")
         table_id = "mock_table_id"
         batch = pa.record_batch([pa.array([1, 2, 3])], names=["a"])
         with responses.RequestsMock() as rsps:
             rsps.add(
                 responses.POST,
                 f"{client.session.base_url}tables/{table_id}/data",
-                status=409
+                status=409,
             )
             with pytest.raises(ApiException) as excinfo:
                 client.append_table_data(table_id, [batch], end_of_data=True)
         assert "Arrow ingestion request was rejected" not in str(excinfo.value)
 
-    def test__append_table_data__unsupported_type_raises(self, client: DataFrameClient, create_table):
+    def test__append_table_data__unsupported_type_raises(
+        self, client: DataFrameClient, create_table
+    ):
         table_id = create_table(basic_table_model)
         with pytest.raises(ValueError, match="Unsupported type"):
             client.append_table_data(table_id, 123)