Remove Arrow support from `obspec.List`

[kylebarron]

Obstore supports returning Arrow RecordBatch from each chunk in obstore.list and returning an Arrow Table from obstore.list_with_delimiter.

I would like this to be an obstore implementation detail instead of a requirement of all obspec implementations.

I had hoped that I would be able to remove the return_arrow keyword from obspec's list methods but still allow obstore's implementation to add the return_arrow keyword as long as it defaults to False and returns a list[ObjectMeta] by default. However it looks like this doesn't pass pylance:

See what I tried in #14

That said, since obspec's list is defined in terms of the Arrow PyCapsule Interface, setting return_arrow=True allows for very generic programming. The return type could be a pandas, Polars, DuckDB, pyarrow, nanoarrow, or arro3 or anything else that supports the protocol. (There is a wrinkle that list requires something that implements the ArrowArray interface, which I'm not sure pandas or polars define, since they only have a concept of a multiple-chunked data structure)

[kylebarron]

With help from chatgpt I got this working! No type errors (at least from pylance)

# _test_list_protocol.py
from __future__ import annotations

from collections.abc import Sequence
from typing import Generic, Protocol, Self, TypedDict, TypeVar

from ._meta import ObjectMeta
from .arrow import ArrowArrayExportable, ArrowStreamExportable

# x Sequence[ObjectMeta],
# ArrowArrayExportable,
# ArrowStreamExportable,

ListChunkType_co = TypeVar(
    "ListChunkType_co",
    covariant=True,
)
"""The data structure used for holding list results.

By default, listing APIs return a `list` of [`ObjectMeta`][obspec.ObjectMeta]. However
for improved performance when listing large buckets, you can pass `return_arrow=True`.
Then an Arrow `RecordBatch` will be returned instead.
"""


class ListResult(TypedDict, Generic[ListChunkType_co]):
    """Result of a `list_with_delimiter` call.

    Includes objects, prefixes (directories) and a token for the next set of results.
    Individual result sets may be limited to 1,000 objects based on the underlying
    object storage's limitations.
    """

    common_prefixes: Sequence[str]
    """Prefixes that are common (like directories)"""

    objects: ListChunkType_co
    """Object metadata for the listing"""


class ListIterator(Protocol[ListChunkType_co]):
    """A stream of [ObjectMeta][obspec.ObjectMeta] that can be polled synchronously."""

    def __iter__(self) -> Self:
        """Return `Self` as an async iterator."""
        ...

    def collect(self) -> ListChunkType_co:
        """Collect all remaining ObjectMeta objects in the stream.

        This ignores the `chunk_size` parameter from the `list` call and collects all
        remaining data into a single chunk.
        """
        ...

    def __next__(self) -> ListChunkType_co:
        """Return the next chunk of ObjectMeta in the stream."""
        ...


class ListStream(Protocol[ListChunkType_co]):
    """A stream of [ObjectMeta][obspec.ObjectMeta] that can be polled asynchronously."""

    def __aiter__(self) -> Self:
        """Return `Self` as an async iterator."""
        ...

    async def collect_async(self) -> ListChunkType_co:
        """Collect all remaining ObjectMeta objects in the stream.

        This ignores the `chunk_size` parameter from the `list` call and collects all
        remaining data into a single chunk.
        """
        ...

    async def __anext__(self) -> ListChunkType_co:
        """Return the next chunk of ObjectMeta in the stream."""
        ...


class List(Protocol):
    def list(
        self,
        prefix: str | None = None,
        *,
        offset: str | None = None,
        chunk_size: int = 50,
    ) -> ListIterator[Sequence[ObjectMeta]]: ...


class ListAsync(Protocol):
    def list_async(
        self,
        prefix: str | None = None,
        *,
        offset: str | None = None,
        chunk_size: int = 50,
    ) -> ListStream[Sequence[ObjectMeta]]: ...


class ListWithDelimiter(Protocol):
    def list_with_delimiter(
        self,
        prefix: str | None = None,
    ) -> ListResult[Sequence[ObjectMeta]]: ...


class ListWithDelimiterAsync(Protocol):
    async def list_with_delimiter_async(
        self,
        prefix: str | None = None,
    ) -> ListResult[Sequence[ObjectMeta]]:
        """Call `list_with_delimiter` asynchronously.

        Refer to the documentation for
        [ListWithDelimiter][obspec.ListWithDelimiter].
        """
        ...

# _test_list_impl.py
from __future__ import annotations

from typing import TYPE_CHECKING, Generic, Literal, Self, Sequence, TypeVar, overload

from arro3.core import RecordBatch, Table

from obspec import ListChunkType_co, ObjectMeta

from ._test_list_protocol import List

if TYPE_CHECKING:
    import obspec


def test_list_arrow_compatible():
    class ListIter(Generic[ListChunkType_co]):
        def __iter__(self) -> Self:
            """Return `Self` as an async iterator."""
            ...

        def collect(self) -> ListChunkType_co:
            """Collect all remaining ObjectMeta objects in the stream.

            This ignores the `chunk_size` parameter from the `list` call and collects all
            remaining data into a single chunk.
            """
            ...

        def __next__(self) -> ListChunkType_co:
            """Return the next chunk of ObjectMeta in the stream."""
            ...

    class ListIterImpl(ListIter[Sequence[ObjectMeta]]):
        def __iter__(self) -> Self:
            """Return `Self` as an async iterator."""
            ...

        def collect(self) -> Sequence[ObjectMeta]:
            """Collect all remaining ObjectMeta objects in the stream.

            This ignores the `chunk_size` parameter from the `list` call and collects all
            remaining data into a single chunk.
            """
            ...

        def __next__(self) -> Sequence[ObjectMeta]:
            """Return the next chunk of ObjectMeta in the stream."""
            ...

    class ObstoreList:
        @overload
        def list(
            self,
            prefix: str | None = None,
            *,
            offset: str | None = None,
            chunk_size: int = 50,
        ) -> ListIter[Sequence[ObjectMeta]]: ...
        @overload
        def list(
            self,
            prefix: str | None = None,
            *,
            offset: str | None = None,
            chunk_size: int = 50,
            return_arrow: Literal[True],
        ) -> ListIter[RecordBatch]: ...
        @overload
        def list(
            self,
            prefix: str | None = None,
            *,
            offset: str | None = None,
            chunk_size: int = 50,
            return_arrow: Literal[False],
        ) -> ListIter[Sequence[ObjectMeta]]: ...
        @overload
        def list(
            self,
            prefix: str | None = None,
            *,
            offset: str | None = None,
            chunk_size: int = 50,
            return_arrow: bool,
        ) -> ListIter[RecordBatch] | ListIter[Sequence[ObjectMeta]]: ...

        def list(
            self,
            prefix: str | None = None,
            *,
            offset: str | None = None,
            chunk_size: int = 50,
            return_arrow: bool = False,
        ) -> ListIter[RecordBatch] | ListIter[Sequence[ObjectMeta]]:
            return ListIterImpl()

    def accepts_obspec_list(provider: List):
        pass

    accepts_obspec_list(ObstoreList())