feat: Add list_async method to align with obspec (#462)

* chore: Test assignability to obspec protocols

* Add list_async alias on store mixin

* remove types-boto3-s3

* restore types-boto3-sts

* fix boto3 type check

Author: kylebarron

obstore/python/obstore/auth/boto3.py

@@ -105,6 +105,8 @@ def __call__(self) -> S3Credential:
         """Fetch credentials."""
         expires_at = datetime.now(timezone.utc) + self.ttl
         frozen_credentials = self.credentials.get_frozen_credentials()
+        assert frozen_credentials.access_key is not None, "Access key is None"
+        assert frozen_credentials.secret_key is not None, "Secret key is None"
         return {
             "access_key_id": frozen_credentials.access_key,
             "secret_access_key": frozen_credentials.secret_key,

obstore/python/obstore/store.py

@@ -325,6 +325,61 @@ def list(
             return_arrow=return_arrow,
         )
 
+    @overload
+    def list_async(
+        self,
+        prefix: str | None = None,
+        *,
+        offset: str | None = None,
+        chunk_size: int = 50,
+        return_arrow: Literal[True],
+    ) -> ListStream[RecordBatch]: ...
+    @overload
+    def list_async(
+        self,
+        prefix: str | None = None,
+        *,
+        offset: str | None = None,
+        chunk_size: int = 50,
+        return_arrow: Literal[False] = False,
+    ) -> ListStream[Sequence[ObjectMeta]]: ...
+    def list_async(
+        self,
+        prefix: str | None = None,
+        *,
+        offset: str | None = None,
+        chunk_size: int = 50,
+        return_arrow: bool = False,
+    ) -> ListStream[RecordBatch] | ListStream[Sequence[ObjectMeta]]:
+        """List all the objects with the given prefix.
+
+        Refer to the documentation for [list][obstore.list].
+
+        !!! note
+
+            This is an alias for `list`, provided to match the `ListAsync` protocol in
+            obspec. There is no difference in functionality between this and the `list`
+            method.
+        """
+        # Splitting these fixes the typing issue with the `return_arrow` parameter, by
+        # converting from a bool to a Literal[True] or Literal[False]
+        if return_arrow:
+            return obs.list(  # type: ignore[call-overload]
+                self,  # type: ignore[arg-type]
+                prefix,
+                offset=offset,
+                chunk_size=chunk_size,
+                return_arrow=return_arrow,
+            )
+
+        return obs.list(  # type: ignore[call-overload]
+            self,  # type: ignore[arg-type]
+            prefix,
+            offset=offset,
+            chunk_size=chunk_size,
+            return_arrow=return_arrow,
+        )
+
     @overload
     def list_with_delimiter(
         self,

pyproject.toml

@@ -12,7 +12,7 @@ dev-dependencies = [
     "aiohttp>=3.11.13",
     "arro3-core>=0.4.2",
     "azure-identity>=1.21.0",
-    "boto3>=1.35.38",
+    "boto3>=1.38.21",
     "fastapi>=0.115.12", # used in example but added here for pyright CI
     "fsspec>=2024.10.0",
     "google-auth>=2.38.0",
@@ -29,7 +29,7 @@ dev-dependencies = [
     "mkdocstrings>=0.27.0",
     "moto[s3,server]>=5.1.1",
     "mypy>=1.15.0",
-    "obspec>=0.1.0b4",
+    "obspec>=0.1.0b5",
     "pip>=24.2",
     "polars>=1.30.0",
     "pyarrow>=17.0.0",
@@ -40,8 +40,8 @@ dev-dependencies = [
     "pytest>=8.3.3",
     "python-dotenv>=1.0.1",
     "ruff>=0.11.0",
-    "tqdm>=4.67.1", # used in example but added here for pyright CI
-    "types-boto3[sts]>=1.36.23",
+    "tqdm>=4.67.1",
+    "types-boto3[s3,sts]>=1.36.23",
 ]
 constraint-dependencies = [
     # ensure lockfile grabs wheels for pyproj for each Python version

tests/obspec/test-store.yml

@@ -1,4 +1,10 @@
-- case: accepts_get
+- case: assignable_to_sync_store
+  parametrized:
+    - store: AzureStore
+    - store: GCSStore
+    - store: LocalStore
+    - store: MemoryStore
+    - store: S3Store
   main: |
     from typing import Protocol
 
@@ -16,8 +22,7 @@
     )
     from typing_extensions import assert_type
 
-    from obstore.store import AzureStore, GCSStore, LocalStore, MemoryStore, S3Store
-
+    from obstore.store import {{ store }}
 
     class SyncStore(
         Copy,
@@ -33,26 +38,54 @@
         Protocol,
     ): ...
 
+    def accepts_sync_store(store: SyncStore) -> None:
+        assert_type(store, SyncStore)
 
-    def memory_store_implements_sync_store(store: MemoryStore) -> None:
-        accepts_sync_store(store)
-
-
-    def azure_store_implements_sync_store(store: AzureStore) -> None:
-        accepts_sync_store(store)
-
-
-    def gcs_store_implements_sync_store(store: GCSStore) -> None:
+    def assignable_to_sync_store(store: {{ store }}) -> None:
         accepts_sync_store(store)
 
+- case: assignable_to_async_store
+  parametrized:
+    - store: AzureStore
+    - store: GCSStore
+    - store: LocalStore
+    - store: MemoryStore
+    - store: S3Store
+  main: |
+    from typing import Protocol
 
-    def local_store_implements_sync_store(store: LocalStore) -> None:
-        accepts_sync_store(store)
-
+    from obspec import (
+        CopyAsync,
+        DeleteAsync,
+        GetAsync,
+        GetRangeAsync,
+        GetRangesAsync,
+        HeadAsync,
+        ListAsync,
+        ListWithDelimiterAsync,
+        PutAsync,
+        RenameAsync,
+    )
+    from typing_extensions import assert_type
 
-    def s3_store_implements_sync_store(store: S3Store) -> None:
-        accepts_sync_store(store)
+    from obstore.store import {{ store }}
+
+    class AsyncStore(
+        CopyAsync,
+        DeleteAsync,
+        GetAsync,
+        GetRangeAsync,
+        GetRangesAsync,
+        HeadAsync,
+        ListAsync,
+        ListWithDelimiterAsync,
+        PutAsync,
+        RenameAsync,
+        Protocol,
+    ): ...
 
+    def accepts_async_store(store: AsyncStore) -> None:
+        assert_type(store, AsyncStore)
 
-    def accepts_sync_store(store: SyncStore) -> None:
-        assert_type(store, SyncStore)
+    def assignable_to_async_store(store: {{ store }}) -> None:
+        accepts_async_store(store)