From 136b92b76197c55434d5ad4b5085cf790867543c Mon Sep 17 00:00:00 2001
From: Alexander Reynolds <alex@theory.shop>
Date: Thu, 12 Dec 2024 16:19:00 -0500
Subject: [PATCH] expand from binaryio -> supportsread

---
 nominal/_utils.py          |  6 +++++-
 nominal/core/_multipart.py |  9 +++++----
 nominal/core/client.py     | 12 ++++++------
 nominal/core/dataset.py    |  5 +++--
 4 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/nominal/_utils.py b/nominal/_utils.py
index f79a11f3..f68e174f 100644
--- a/nominal/_utils.py
+++ b/nominal/_utils.py
@@ -3,7 +3,7 @@
 import logging
 import os
 from contextlib import contextmanager
-from typing import Any, BinaryIO, Callable, Iterator, TypeVar
+from typing import Any, BinaryIO, Callable, Iterator, Protocol, TypeVar
 
 from typing_extensions import ParamSpec
 
@@ -64,3 +64,7 @@ def wrapper(*args: Param.args, **kwargs: Param.kwargs) -> T:
         return wrapper
 
     return _deprecate_keyword_argument_decorator
+
+
+class HasBinaryRead(Protocol):
+    def read(self, length: int = ..., /) -> bytes: ...
diff --git a/nominal/core/_multipart.py b/nominal/core/_multipart.py
index 68165c1f..eed4453b 100644
--- a/nominal/core/_multipart.py
+++ b/nominal/core/_multipart.py
@@ -6,11 +6,12 @@
 import urllib.parse
 from functools import partial
 from queue import Queue
-from typing import BinaryIO, Iterable
+from typing import Iterable
 
 import requests
 
 from nominal._api.scout_service_api import ingest_api, upload_api
+from nominal._utils import HasBinaryRead
 from nominal.core.filetype import FileType
 from nominal.exceptions import NominalMultipartUploadFailed
 
@@ -49,14 +50,14 @@ def _sign_and_upload_part_job(
         q.task_done()
 
 
-def _iter_chunks(f: BinaryIO, chunk_size: int) -> Iterable[bytes]:
+def _iter_chunks(f: HasBinaryRead, chunk_size: int) -> Iterable[bytes]:
     while (data := f.read(chunk_size)) != b"":
         yield data
 
 
 def put_multipart_upload(
     auth_header: str,
-    f: BinaryIO,
+    f: HasBinaryRead,
     filename: str,
     mimetype: str,
     upload_client: upload_api.UploadService,
@@ -138,7 +139,7 @@ def put_multipart_upload(
 
 def upload_multipart_io(
     auth_header: str,
-    f: BinaryIO,
+    f: HasBinaryRead,
     name: str,
     file_type: FileType,
     upload_client: upload_api.UploadService,
diff --git a/nominal/core/client.py b/nominal/core/client.py
index 12457421..a8af0475 100644
--- a/nominal/core/client.py
+++ b/nominal/core/client.py
@@ -6,7 +6,7 @@
 from datetime import datetime
 from io import BytesIO, TextIOBase, TextIOWrapper
 from pathlib import Path
-from typing import BinaryIO, Iterable, Mapping, Sequence
+from typing import Iterable, Mapping, Sequence
 
 import certifi
 from conjure_python_client import ServiceConfiguration, SslConfiguration
@@ -28,7 +28,7 @@
     storage_datasource_api,
     timeseries_logicalseries_api,
 )
-from nominal._utils import deprecate_keyword_argument
+from nominal._utils import HasBinaryRead, deprecate_keyword_argument
 from nominal.core._clientsbunch import ClientsBunch
 from nominal.core._conjure_utils import _available_units, _build_unit_update
 from nominal.core._multipart import upload_multipart_file, upload_multipart_io
@@ -284,7 +284,7 @@ def create_mcap_dataset(
 
     def create_dataset_from_io(
         self,
-        dataset: BinaryIO,
+        dataset: HasBinaryRead,
         name: str,
         timestamp_column: str,
         timestamp_type: _AnyTimestampType,
@@ -367,7 +367,7 @@ def create_video(
 
     def create_video_from_io(
         self,
-        video: BinaryIO,
+        video: HasBinaryRead,
         name: str,
         start: datetime | IntegralNanosecondsUTC | None = None,
         frame_timestamps: Sequence[IntegralNanosecondsUTC] | None = None,
@@ -567,7 +567,7 @@ def checklist_builder(
 
     def create_attachment_from_io(
         self,
-        attachment: BinaryIO,
+        attachment: HasBinaryRead,
         name: str,
         file_type: tuple[str, str] | FileType = FileTypes.BINARY,
         description: str | None = None,
@@ -689,7 +689,7 @@ def get_connection(self, rid: str) -> Connection:
 
     def create_video_from_mcap_io(
         self,
-        mcap: BinaryIO,
+        mcap: HasBinaryRead,
         topic: str,
         name: str,
         description: str | None = None,
diff --git a/nominal/core/dataset.py b/nominal/core/dataset.py
index 107a3822..d524fa39 100644
--- a/nominal/core/dataset.py
+++ b/nominal/core/dataset.py
@@ -7,7 +7,7 @@
 from io import TextIOBase
 from pathlib import Path
 from types import MappingProxyType
-from typing import BinaryIO, Iterable, Mapping, Protocol, Sequence
+from typing import Iterable, Mapping, Protocol, Sequence
 
 import pandas as pd
 from typing_extensions import Self
@@ -23,6 +23,7 @@
     timeseries_logicalseries_api,
     upload_api,
 )
+from nominal._utils import HasBinaryRead
 from nominal.core._clientsbunch import HasAuthHeader
 from nominal.core._conjure_utils import _available_units, _build_unit_update
 from nominal.core._multipart import upload_multipart_io
@@ -173,7 +174,7 @@ def add_data_to_dataset(self, path: Path | str, timestamp_column: str, timestamp
 
     def add_to_dataset_from_io(
         self,
-        dataset: BinaryIO,
+        dataset: HasBinaryRead,
         timestamp_column: str,
         timestamp_type: _AnyTimestampType,
         file_type: tuple[str, str] | FileType = FileTypes.CSV,