earthscope-sdk 1.1.0__py3-none-any.whl → 1.2.0b0__py3-none-any.whl

earthscope_sdk/client/discovery/_base.py

@@ -0,0 +1,303 @@
+import datetime as dt
+from typing import Optional, Union
+
+from earthscope_sdk.client.discovery.models import (
+    ListNetworkDatasourcesResult,
+    ListSessionDatasourcesResult,
+    ListStationDatasourcesResult,
+    ListStreamDatasourcesResult,
+    NetworkDatasource,
+    Page,
+    SessionDatasource,
+    StationDatasource,
+    StreamDatasource,
+    StreamType,
+)
+from earthscope_sdk.common.service import SdkService
+from earthscope_sdk.util._itertools import to_list
+from earthscope_sdk.util._types import ListOrItem
+
+
+class DiscoveryBaseService(SdkService):
+    """
+    Discovery service functionality
+    """
+
+    async def _list_network_datasources(
+        self,
+        *,
+        network_name: Optional[ListOrItem[str]] = None,
+        network_edid: Optional[ListOrItem[str]] = None,
+        with_edid_only=False,
+        limit=50,
+        offset=0,
+    ) -> Union[Page[NetworkDatasource], Page[str]]:
+        """
+        List network datasources
+
+        Args:
+            network_name: The name(s) of the network to list datasources for.
+            network_edid: The EDID(s) of the network to list datasources for.
+            with_edid_only: Whether to return only the EDIDs of the datasources.
+            limit: The maximum number of datasources to return.
+            offset: The offset to start the list from.
+
+        Returns:
+            A list of network datasources.
+        """
+        params = {
+            "with_edid_only": with_edid_only,
+            "limit": limit,
+            "offset": offset,
+        }
+
+        if network_name:
+            params["network_name"] = to_list(network_name)
+
+        if network_edid:
+            params["network_edid"] = to_list(network_edid)
+
+        req = self.ctx.httpx_client.build_request(
+            method="GET",
+            url=f"{self.resources.api_url}beta/discover/datasource/network",
+            params=params,
+        )
+
+        resp = await self._send_with_retries(req)
+
+        return ListNetworkDatasourcesResult.validate_json(resp.content)
+
+    async def _list_station_datasources(
+        self,
+        *,
+        network_name: Optional[ListOrItem[str]] = None,
+        network_edid: Optional[ListOrItem[str]] = None,
+        station_name: Optional[ListOrItem[str]] = None,
+        station_edid: Optional[ListOrItem[str]] = None,
+        with_edid_only=False,
+        with_parent_edids=False,
+        limit=50,
+        offset=0,
+    ) -> Union[Page[StationDatasource], Page[str]]:
+        """
+        List station datasources
+
+        Args:
+            network_name: The name(s) of the network to list datasources for.
+            network_edid: The EDID(s) of the network to list datasources for.
+            station_name: The name(s) of the station to list datasources for.
+            station_edid: The EDID(s) of the station to list datasources for.
+            with_edid_only: Whether to return only the EDIDs of the datasources.
+            with_parent_edids: Whether to return the parent EDIDs of the datasources.
+            limit: The maximum number of datasources to return.
+            offset: The offset to start the list from.
+
+        Returns:
+            A list of station datasources.
+        """
+        params = {
+            "with_edid_only": with_edid_only,
+            "with_parent_edids": with_parent_edids,
+            "limit": limit,
+            "offset": offset,
+        }
+
+        if network_name:
+            params["network_name"] = to_list(network_name)
+
+        if network_edid:
+            params["network_edid"] = to_list(network_edid)
+
+        if station_name:
+            params["station_name"] = to_list(station_name)
+
+        if station_edid:
+            params["station_edid"] = to_list(station_edid)
+
+        req = self.ctx.httpx_client.build_request(
+            method="GET",
+            url=f"{self.resources.api_url}beta/discover/datasource/station",
+            params=params,
+        )
+
+        resp = await self._send_with_retries(req)
+
+        return ListStationDatasourcesResult.validate_json(resp.content)
+
+    async def _list_session_datasources(
+        self,
+        *,
+        network_name: Optional[ListOrItem[str]] = None,
+        network_edid: Optional[ListOrItem[str]] = None,
+        station_name: Optional[ListOrItem[str]] = None,
+        station_edid: Optional[ListOrItem[str]] = None,
+        session_name: Optional[ListOrItem[str]] = None,
+        session_edid: Optional[ListOrItem[str]] = None,
+        sample_interval: Optional[dt.timedelta] = None,
+        roll: Optional[dt.timedelta] = None,
+        with_edid_only=False,
+        with_parents=False,
+        with_parent_edids=False,
+        limit=50,
+        offset=0,
+    ) -> Union[Page[SessionDatasource], Page[str]]:
+        """
+        List session datasources
+
+        Args:
+            network_name: The name(s) of the network to list datasources for.
+            network_edid: The EDID(s) of the network to list datasources for.
+            station_name: The name(s) of the station to list datasources for.
+            station_edid: The EDID(s) of the station to list datasources for.
+            session_name: The name(s) of the session to list datasources for.
+            session_edid: The EDID(s) of the session to list datasources for.
+            sample_interval: The sample interval to list datasources for.
+            roll: The roll to list datasources for.
+            with_edid_only: Whether to return only the EDIDs of the datasources.
+            with_parents: Whether to return the parent datasources.
+            with_parent_edids: Whether to return the parent EDIDs of the datasources.
+            limit: The maximum number of datasources to return.
+            offset: The offset to start the list from.
+
+        Returns:
+            A list of session datasources.
+        """
+        params = {
+            "with_edid_only": with_edid_only,
+            "with_parents": with_parents,
+            "with_parent_edids": with_parent_edids,
+            "limit": limit,
+            "offset": offset,
+        }
+
+        if network_name:
+            params["network_name"] = to_list(network_name)
+
+        if network_edid:
+            params["network_edid"] = to_list(network_edid)
+
+        if station_name:
+            params["station_name"] = to_list(station_name)
+
+        if station_edid:
+            params["station_edid"] = to_list(station_edid)
+
+        if session_name:
+            params["session_name"] = to_list(session_name)
+
+        if session_edid:
+            params["session_edid"] = to_list(session_edid)
+
+        if sample_interval:
+            # Convert to milliseconds
+            params["sample_interval"] = 1000 * sample_interval.total_seconds()
+
+        if roll:
+            # Convert to seconds
+            params["roll"] = roll.total_seconds()
+
+        req = self.ctx.httpx_client.build_request(
+            method="GET",
+            url=f"{self.resources.api_url}beta/discover/datasource/session",
+            params=params,
+        )
+
+        resp = await self._send_with_retries(req)
+
+        return ListSessionDatasourcesResult.validate_json(resp.content)
+
+    async def _list_stream_datasources(
+        self,
+        *,
+        network_name: Optional[ListOrItem[str]] = None,
+        network_edid: Optional[ListOrItem[str]] = None,
+        station_name: Optional[ListOrItem[str]] = None,
+        station_edid: Optional[ListOrItem[str]] = None,
+        stream_name: Optional[ListOrItem[str]] = None,
+        stream_edid: Optional[ListOrItem[str]] = None,
+        stream_type: Optional[StreamType] = None,
+        facility: Optional[str] = None,
+        software: Optional[str] = None,
+        label: Optional[str] = None,
+        sample_interval: Optional[dt.timedelta] = None,
+        with_edid_only=False,
+        with_parents=False,
+        with_parent_edids=False,
+        limit=50,
+        offset=0,
+    ) -> Union[Page[StreamDatasource], Page[str]]:
+        """
+        List stream datasources
+
+        Args:
+            network_name: The name(s) of the network to list datasources for.
+            network_edid: The EDID(s) of the network to list datasources for.
+            station_name: The name(s) of the station to list datasources for.
+            station_edid: The EDID(s) of the station to list datasources for.
+            stream_name: The name(s) of the stream to list datasources for.
+            stream_edid: The EDID(s) of the stream to list datasources for.
+            stream_type: The type of stream to list datasources for.
+            facility: The facility to list datasources for.
+            software: The software to list datasources for.
+            label: The label to list datasources for.
+            sample_interval: The sample interval to list datasources for.
+            with_edid_only: Whether to return only the EDIDs of the datasources.
+            with_parents: Whether to return the parent datasources.
+            with_parent_edids: Whether to return the parent EDIDs of the datasources.
+            limit: The maximum number of datasources to return.
+            offset: The offset to start the list from.
+
+        Returns:
+            A list of stream datasources.
+        """
+        params = {
+            "with_edid_only": with_edid_only,
+            "with_parents": with_parents,
+            "with_parent_edids": with_parent_edids,
+            "limit": limit,
+            "offset": offset,
+        }
+
+        if network_name:
+            params["network_name"] = to_list(network_name)
+
+        if network_edid:
+            params["network_edid"] = to_list(network_edid)
+
+        if station_name:
+            params["station_name"] = to_list(station_name)
+
+        if station_edid:
+            params["station_edid"] = to_list(station_edid)
+
+        if stream_name:
+            params["stream_name"] = to_list(stream_name)
+
+        if stream_edid:
+            params["stream_edid"] = to_list(stream_edid)
+
+        if stream_type is not None:
+            params["stream_type"] = stream_type.value
+
+        if facility is not None:
+            params["facility"] = facility
+
+        if software is not None:
+            params["software"] = software
+
+        if label is not None:
+            params["label"] = label
+
+        if sample_interval is not None:
+            # Convert to milliseconds
+            params["sample_interval"] = 1000 * sample_interval.total_seconds()
+
+        req = self.ctx.httpx_client.build_request(
+            method="GET",
+            url=f"{self.resources.api_url}beta/discover/datasource/stream",
+            params=params,
+        )
+
+        resp = await self._send_with_retries(req)
+
+        return ListStreamDatasourcesResult.validate_json(resp.content)

earthscope_sdk/client/discovery/_service.py

@@ -0,0 +1,209 @@
+import datetime as dt
+from functools import partial
+from typing import Any, AsyncIterator, Coroutine, Optional, TypeVar, Union
+
+from earthscope_sdk.client.discovery._base import DiscoveryBaseService
+from earthscope_sdk.client.discovery.models import (
+    NetworkDatasource,
+    Page,
+    SessionDatasource,
+    StationDatasource,
+    StreamDatasource,
+    StreamType,
+)
+from earthscope_sdk.common.context import SdkContext
+from earthscope_sdk.util._types import ListOrItem
+
+P = TypeVar("P")
+
+
+class _DiscoveryService(DiscoveryBaseService):
+    """
+    L2 discovery service functionality
+    """
+
+    async def _iter_pages(
+        self,
+        fn: Coroutine[Any, Any, Page[P]],
+        limit: int,
+    ) -> AsyncIterator[Page[P]]:
+        """
+        Iterate over pages of results from a function.
+
+        Args:
+            fn: The function to call to get a page of results.
+            limit: The maximum number of results to return.
+
+        Returns:
+            An iterator over the pages of results.
+        """
+        offset = 0
+        page_size = 100
+        has_next = True
+        result_ct = 0
+
+        while has_next and result_ct < limit:
+            page_size = min(page_size, limit - result_ct)
+            page: Page[P] = await fn(offset=offset, limit=page_size)
+            yield page
+            has_next = page.has_next
+            offset += page_size
+            result_ct += len(page.items)
+
+    async def _load_all_pages(
+        self,
+        fn: Coroutine[Any, Any, Page[P]],
+        limit: int,
+    ) -> list[P]:
+        """
+        Load all pages of results from a function.
+        """
+        results: list[P] = []
+        async for p in self._iter_pages(fn=fn, limit=limit):
+            results.extend(p.items)
+
+        return results
+
+    async def _list_network_datasources(
+        self,
+        *,
+        network_name: Optional[ListOrItem[str]] = None,
+        network_edid: Optional[ListOrItem[str]] = None,
+        with_edid_only=False,
+        limit=1000,
+    ) -> Union[list[NetworkDatasource], list[str]]:
+        return await self._load_all_pages(
+            fn=partial(
+                super()._list_network_datasources,
+                network_name=network_name,
+                network_edid=network_edid,
+                with_edid_only=with_edid_only,
+            ),
+            limit=limit,
+        )
+
+    async def _list_station_datasources(
+        self,
+        *,
+        network_name: Optional[ListOrItem[str]] = None,
+        network_edid: Optional[ListOrItem[str]] = None,
+        station_name: Optional[ListOrItem[str]] = None,
+        station_edid: Optional[ListOrItem[str]] = None,
+        with_edid_only=False,
+        with_parent_edids=False,
+        limit=1000,
+    ) -> Union[list[StationDatasource], list[str]]:
+        return await self._load_all_pages(
+            fn=partial(
+                super()._list_station_datasources,
+                network_name=network_name,
+                network_edid=network_edid,
+                station_name=station_name,
+                station_edid=station_edid,
+                with_edid_only=with_edid_only,
+                with_parent_edids=with_parent_edids,
+            ),
+            limit=limit,
+        )
+
+    async def _list_session_datasources(
+        self,
+        *,
+        network_name: Optional[ListOrItem[str]] = None,
+        network_edid: Optional[ListOrItem[str]] = None,
+        station_name: Optional[ListOrItem[str]] = None,
+        station_edid: Optional[ListOrItem[str]] = None,
+        session_name: Optional[ListOrItem[str]] = None,
+        session_edid: Optional[ListOrItem[str]] = None,
+        sample_interval: Optional[dt.timedelta] = None,
+        roll: Optional[dt.timedelta] = None,
+        with_edid_only=False,
+        with_parents=False,
+        with_parent_edids=False,
+        limit=1000,
+    ) -> Union[list[SessionDatasource], list[str]]:
+        return await self._load_all_pages(
+            fn=partial(
+                super()._list_session_datasources,
+                network_name=network_name,
+                network_edid=network_edid,
+                station_name=station_name,
+                station_edid=station_edid,
+                session_name=session_name,
+                session_edid=session_edid,
+                sample_interval=sample_interval,
+                roll=roll,
+                with_edid_only=with_edid_only,
+                with_parents=with_parents,
+                with_parent_edids=with_parent_edids,
+            ),
+            limit=limit,
+        )
+
+    async def _list_stream_datasources(
+        self,
+        *,
+        network_name: Optional[ListOrItem[str]] = None,
+        network_edid: Optional[ListOrItem[str]] = None,
+        station_name: Optional[ListOrItem[str]] = None,
+        station_edid: Optional[ListOrItem[str]] = None,
+        stream_name: Optional[ListOrItem[str]] = None,
+        stream_edid: Optional[ListOrItem[str]] = None,
+        stream_type: Optional[StreamType] = None,
+        facility: Optional[str] = None,
+        software: Optional[str] = None,
+        label: Optional[str] = None,
+        sample_interval: Optional[dt.timedelta] = None,
+        with_edid_only=False,
+        with_parents=False,
+        with_parent_edids=False,
+        limit=1000,
+    ) -> Union[list[StreamDatasource], list[str]]:
+        return await self._load_all_pages(
+            fn=partial(
+                super()._list_stream_datasources,
+                network_name=network_name,
+                network_edid=network_edid,
+                station_name=station_name,
+                station_edid=station_edid,
+                stream_name=stream_name,
+                stream_edid=stream_edid,
+                stream_type=stream_type,
+                facility=facility,
+                software=software,
+                label=label,
+                sample_interval=sample_interval,
+                with_edid_only=with_edid_only,
+                with_parents=with_parents,
+                with_parent_edids=with_parent_edids,
+            ),
+            limit=limit,
+        )
+
+
+class AsyncDiscoveryService(_DiscoveryService):
+    """
+    Discovery service functionality
+    """
+
+    def __init__(self, ctx: SdkContext):
+        super().__init__(ctx)
+
+        self.list_network_datasources = self._list_network_datasources
+        self.list_station_datasources = self._list_station_datasources
+        self.list_session_datasources = self._list_session_datasources
+        self.list_stream_datasources = self._list_stream_datasources
+
+
+class DiscoveryService(_DiscoveryService):
+    """
+    Discovery service functionality
+    """
+
+    def __init__(self, ctx: SdkContext):
+        super().__init__(ctx)
+
+        self.list_network_datasources = ctx.syncify(self._list_network_datasources)
+        self.list_station_datasources = ctx.syncify(self._list_station_datasources)
+        self.list_session_datasources = ctx.syncify(self._list_session_datasources)
+        self.list_stream_datasources = ctx.syncify(self._list_stream_datasources)

earthscope_sdk/client/discovery/models.py

@@ -0,0 +1,144 @@
+from datetime import timedelta
+from enum import Enum
+from typing import Annotated, Any, Generic, Iterable, Optional, TypeVar, Union
+
+from pydantic import BaseModel, BeforeValidator, TypeAdapter
+
+from earthscope_sdk.util._itertools import to_set
+
+
+def _coerce_timedelta_ms(v: Union[int, float, timedelta, str]) -> timedelta:
+    if isinstance(v, (int, float)):
+        return timedelta(milliseconds=v)
+
+    # fallback to Pydantic's timedelta parser
+    return v
+
+
+P = TypeVar("P")
+
+
+class Page(BaseModel, Generic[P]):
+    has_next: bool
+    offset: int
+    limit: int
+    items: list[P]
+    total: Optional[int] = None
+
+
+class DatasourceBaseModel(BaseModel):
+    edid: str
+    names: dict[str, str]
+    description: Optional[str] = None
+
+    def to_arrow_columns(
+        self,
+        *,
+        fields: Union[list[str], str] = ["edid", "names"],
+        namespaces: Union[list[str], str] = [],
+    ) -> dict[str, Any]:
+        """
+        Convert the datasource model to a dictionary suitable for use in an Arrow table.
+        """
+        result = {}
+        namespaces = to_set(namespaces)
+        fields = to_set(fields)
+
+        # Add names to fields if namespaces are requested
+        if namespaces:
+            fields.add("names")
+
+        for field in fields:
+            if field != "names":
+                result[field] = getattr(self, field)
+                continue
+
+            # Explode names to own columns
+            if not namespaces:
+                names = {k.lower(): v for k, v in self.names.items()}
+            else:
+                names = {
+                    k_lower: v
+                    for k, v in self.names.items()
+                    if (k_lower := k.lower()) in namespaces
+                }
+
+            result.update(names)
+
+        return result
+
+
+class NetworkDatasource(DatasourceBaseModel): ...
+
+
+ListNetworkDatasourcesResult = TypeAdapter(Union[Page[str], Page[NetworkDatasource]])
+
+
+class StationDatasource(DatasourceBaseModel):
+    network_edids: Optional[list[str]] = None
+    networks: Optional[list[NetworkDatasource]] = None
+
+
+ListStationDatasourcesResult = TypeAdapter(Union[Page[str], Page[StationDatasource]])
+
+
+class _StationDatasourceMember(DatasourceBaseModel):
+    station_edid: Optional[str] = None
+    station: Optional[StationDatasource] = None
+
+    def to_arrow_columns(
+        self,
+        *,
+        fields: list[str] = ["edid", "names"],
+        namespaces: Optional[list[str]] = None,
+    ) -> dict[str, Any]:
+        result = super().to_arrow_columns(fields=fields, namespaces=namespaces)
+        if self.station:
+            parent_columns = self.station.to_arrow_columns(
+                fields=["names"],
+                namespaces=namespaces,
+            )
+            result.update(parent_columns)
+
+        return result
+
+
+class SessionDatasource(_StationDatasourceMember):
+    sample_interval: Annotated[timedelta, BeforeValidator(_coerce_timedelta_ms)]
+    """
+    Session sample interval.
+    """
+
+    roll: timedelta  # already in seconds
+    """
+    Session file roll cadence.
+    """
+
+
+ListSessionDatasourcesResult = TypeAdapter(Union[Page[str], Page[SessionDatasource]])
+
+
+class StreamType(Enum):
+    GNSS_RAW = "gnss_raw"
+    GNSS_PPP = "gnss_ppp"
+
+
+class StreamDatasource(_StationDatasourceMember):
+    stream_type: StreamType
+    facility: str
+    software: str
+    label: str
+    sample_interval: Annotated[timedelta, BeforeValidator(_coerce_timedelta_ms)]
+    """
+    Stream sample interval.
+    """
+
+    def to_arrow_columns(
+        self,
+        *,
+        fields: Iterable[str] = ["edid", "names", "facility", "software", "label"],
+    ) -> dict[str, Any]:
+        return super().to_arrow_columns(fields=fields)
+
+
+ListStreamDatasourcesResult = TypeAdapter(Union[Page[str], Page[StreamDatasource]])