floorplan-api 0.5.0__py3-none-any.whl

floorplan_api/async_client.py

@@ -0,0 +1,400 @@
+"""asyncio client for the Floor Plan API (built on ``httpx``).
+
+Install with the extra::
+
+    pip install "floorplan-api[async]"
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import Mapping
+from typing import Any
+
+try:
+    import httpx
+except ImportError as exc:  # pragma: no cover - exercised only without the extra
+    raise ImportError(
+        "AsyncClient requires httpx. Install it with: pip install 'floorplan-api[async]'"
+    ) from exc
+
+from floorplan_api._input import (
+    ImageInput,
+    ResolvedInput,
+    normalize_content_type,
+    prepare_input,
+)
+from floorplan_api._transport import (
+    DEFAULT_MAX_RETRIES,
+    DEFAULT_PRESIGN_THRESHOLD,
+    DEFAULT_RETRY_BACKOFF,
+    DEFAULT_TIMEOUT,
+    build_headers,
+    build_url,
+    decode_json,
+    error_from_response,
+    job_from_submit,
+    mask_from_response,
+    parse_retry_after,
+    resolve_credentials,
+    retry_delay,
+    should_retry,
+)
+from floorplan_api.exceptions import (
+    ConnectionError as ApiConnectionError,
+)
+from floorplan_api.exceptions import (
+    InvalidRequestError,
+    ServerError,
+)
+from floorplan_api.exceptions import (
+    TimeoutError as ApiTimeoutError,
+)
+from floorplan_api.models import Job, MaskBytes
+
+__all__ = ["AsyncClient"]
+
+
+class AsyncClient:
+    """asyncio Floor Plan API client. Same surface as :class:`~floorplan_api.Client`,
+    every method is a coroutine.
+
+    Use it as an async context manager so the connection pool is closed::
+
+        async with AsyncClient(api_key="fp_test_...") as client:
+            mask = await client.extract("plan.pdf")
+
+    Args:
+        api_key: Your API key, or ``FLOORPLAN_API_KEY`` from the environment.
+        base_url: API base URL, or ``FLOORPLAN_BASE_URL``; defaults to
+            ``https://api.floorplanapi.com``.
+        timeout: Per-request timeout in seconds. Default: 60.
+        max_retries: Retries for network failures, 429s, and transient 5xx.
+            Default: 3.
+        retry_backoff: Base delay for exponential backoff. Default: 0.5.
+        client: An existing :class:`httpx.AsyncClient` to reuse (proxies,
+            custom transport, HTTP/2). If omitted, one is created.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        *,
+        base_url: str | None = None,
+        timeout: float = DEFAULT_TIMEOUT,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+        retry_backoff: float = DEFAULT_RETRY_BACKOFF,
+        client: httpx.AsyncClient | None = None,
+    ) -> None:
+        self.api_key, self.base_url = resolve_credentials(api_key, base_url)
+        self.timeout = timeout
+        self.max_retries = max(0, int(max_retries))
+        self.retry_backoff = max(0.0, float(retry_backoff))
+        self._client = client or httpx.AsyncClient()
+        self._owns_client = client is None
+
+    # ----- public API -----
+
+    async def extract(
+        self,
+        image: ImageInput | None = None,
+        *,
+        upload_key: str | None = None,
+        page: int | None = None,
+        timeout: float | None = None,
+    ) -> MaskBytes:
+        """Run wall segmentation and return the mask PNG. See
+        :py:meth:`floorplan_api.Client.extract`."""
+
+        return await self._submit_for_mask("/v1/extract", image, upload_key, page, timeout)
+
+    async def analyze(
+        self,
+        image: ImageInput | None = None,
+        *,
+        upload_key: str | None = None,
+        page: int | None = None,
+        timeout: float | None = None,
+    ) -> MaskBytes:
+        """Same wall mask as :py:meth:`extract`. See
+        :py:meth:`floorplan_api.Client.analyze`."""
+
+        return await self._submit_for_mask("/v1/analyze", image, upload_key, page, timeout)
+
+    async def analyze_async(
+        self,
+        image: ImageInput | None = None,
+        *,
+        upload_key: str | None = None,
+        page: int | None = None,
+        timeout: float | None = None,
+    ) -> Job:
+        """Submit a job and return its :class:`~floorplan_api.Job` immediately."""
+
+        resolved, files = _multipart(image, upload_key, page, {"async": "true"})
+        try:
+            response = await self._request(
+                "POST", "/v1/analyze", files=files, timeout=timeout, rewind=resolved
+            )
+        finally:
+            if resolved is not None:
+                resolved.close()
+        return job_from_submit(response.json())
+
+    async def get_job(self, job_id: str, *, timeout: float | None = None) -> Job:
+        """Fetch the latest status of an asynchronous job."""
+
+        response = await self._request("GET", f"/v1/status/{job_id}", timeout=timeout)
+        return Job.from_dict(response.json())
+
+    async def wait_for_job(
+        self,
+        job_id: str,
+        *,
+        poll_interval: float = 1.0,
+        timeout: float = 300.0,
+    ) -> Job:
+        """Poll a job until it reaches a terminal status. See
+        :py:meth:`floorplan_api.Client.wait_for_job`."""
+
+        deadline = time.monotonic() + timeout
+        while True:
+            job = await self.get_job(job_id)
+            if job.is_terminal:
+                return job
+            if time.monotonic() > deadline:
+                raise ApiTimeoutError(
+                    f"Job {job_id} did not finish within {timeout} seconds.", job_id=job_id
+                )
+            await asyncio.sleep(poll_interval)
+
+    async def download_mask(self, job_id: str, *, timeout: float | None = None) -> MaskBytes:
+        """Download the wall mask PNG for a completed job."""
+
+        response = await self._request("GET", f"/v1/jobs/{job_id}/mask", timeout=timeout)
+        return mask_from_response(response.content, response.headers)
+
+    async def create_upload(
+        self,
+        *,
+        content_type: str = "image/png",
+        timeout: float | None = None,
+    ) -> dict[str, Any]:
+        """Reserve a presigned upload slot. See
+        :py:meth:`floorplan_api.Client.create_upload`."""
+
+        normalized = normalize_content_type(content_type)
+        if normalized is None:
+            raise InvalidRequestError(
+                f"Unsupported content_type {content_type!r}. Use image/png, image/jpeg, "
+                "image/webp or application/pdf."
+            )
+        response = await self._request(
+            "POST", "/v1/uploads", json_body={"content_type": normalized}, timeout=timeout
+        )
+        payload = response.json()
+        if not isinstance(payload, dict):
+            raise ServerError("Unexpected response from /v1/uploads.", status_code=None)
+        return payload
+
+    async def upload(
+        self,
+        image: ImageInput,
+        *,
+        page: int | None = None,
+        content_type: str | None = None,
+        timeout: float | None = None,
+    ) -> str:
+        """PUT the file directly to object storage and return its storage key.
+        See :py:meth:`floorplan_api.Client.upload`."""
+
+        resolved = prepare_input(image, page)
+        return await self._upload_prepared(resolved, content_type, timeout)
+
+    async def _upload_prepared(
+        self,
+        resolved: ResolvedInput,
+        content_type: str | None,
+        timeout: float | None,
+    ) -> str:
+        try:
+            ct = normalize_content_type(content_type) or resolved.content_type
+            presigned = await self.create_upload(content_type=ct, timeout=timeout)
+            request_timeout = self.timeout if timeout is None else timeout
+            try:
+                put = await self._client.put(
+                    presigned["upload_url"],
+                    content=resolved.file.read(),
+                    headers={"Content-Type": ct},
+                    timeout=request_timeout,
+                )
+            except httpx.TimeoutException as exc:
+                raise ApiTimeoutError(str(exc)) from exc
+            except httpx.TransportError as exc:
+                raise ApiConnectionError(str(exc)) from exc
+            if put.status_code >= 400:
+                raise ServerError(
+                    f"Upload PUT failed ({put.status_code}): {put.text[:200]}",
+                    status_code=put.status_code,
+                )
+        finally:
+            resolved.close()
+        key = presigned.get("key")
+        if not isinstance(key, str) or not key:
+            raise ServerError("Upload response missing `key`.", status_code=None)
+        return key
+
+    async def upload_then_extract(
+        self,
+        image: ImageInput,
+        *,
+        page: int | None = None,
+        threshold_bytes: int = DEFAULT_PRESIGN_THRESHOLD,
+        timeout: float | None = None,
+    ) -> MaskBytes:
+        """Extract walls, switching to the presigned flow for large files.
+        See :py:meth:`floorplan_api.Client.upload_then_extract`."""
+
+        resolved = prepare_input(image, page)
+        size = resolved.size()
+        if size is not None and size < threshold_bytes:
+            return await self._submit_prepared("/v1/extract", resolved, {}, timeout)
+        key = await self._upload_prepared(resolved, None, timeout)
+        return await self._submit_files(
+            "/v1/extract",
+            None,
+            {"upload_key": (None, key)},
+            timeout,
+            page_size_pt=resolved.page_size_pt,
+        )
+
+    async def aclose(self) -> None:
+        """Close the underlying connection pool (only if this client created it)."""
+
+        if self._owns_client:
+            await self._client.aclose()
+
+    async def __aenter__(self) -> AsyncClient:
+        return self
+
+    async def __aexit__(self, *exc: Any) -> None:
+        await self.aclose()
+
+    # ----- internals -----
+
+    async def _submit_for_mask(
+        self,
+        path: str,
+        image: ImageInput | None,
+        upload_key: str | None,
+        page: int | None,
+        timeout: float | None,
+    ) -> MaskBytes:
+        resolved, files = _multipart(image, upload_key, page, {})
+        return await self._submit_files(path, resolved, files, timeout)
+
+    async def _submit_prepared(
+        self,
+        path: str,
+        resolved: ResolvedInput,
+        extra: Mapping[str, str],
+        timeout: float | None,
+    ) -> MaskBytes:
+        files: dict[str, Any] = {name: (None, value) for name, value in extra.items()}
+        files["image"] = resolved.part
+        return await self._submit_files(
+            path, resolved, files, timeout, page_size_pt=resolved.page_size_pt
+        )
+
+    async def _submit_files(
+        self,
+        path: str,
+        resolved: ResolvedInput | None,
+        files: dict[str, Any],
+        timeout: float | None,
+        page_size_pt: tuple[float, float] | None = None,
+    ) -> MaskBytes:
+        if resolved is not None and page_size_pt is None:
+            page_size_pt = resolved.page_size_pt
+        try:
+            response = await self._request(
+                "POST", path, files=files, timeout=timeout, rewind=resolved
+            )
+        finally:
+            if resolved is not None:
+                resolved.close()
+        return mask_from_response(response.content, response.headers, page_size_pt)
+
+    async def _request(
+        self,
+        method: str,
+        path: str,
+        *,
+        files: Mapping[str, Any] | None = None,
+        json_body: Any | None = None,
+        timeout: float | None = None,
+        rewind: ResolvedInput | None = None,
+    ) -> httpx.Response:
+        url = build_url(self.base_url, path)
+        request_timeout = self.timeout if timeout is None else timeout
+        attempt = 0
+        while True:
+            if attempt and rewind is not None:
+                rewind.rewind()
+            try:
+                response = await self._client.request(
+                    method,
+                    url,
+                    headers=build_headers(self.api_key),
+                    files=files,
+                    json=json_body,
+                    timeout=request_timeout,
+                )
+            except httpx.TimeoutException as exc:
+                if attempt < self.max_retries:
+                    await asyncio.sleep(retry_delay(attempt, self.retry_backoff))
+                    attempt += 1
+                    continue
+                raise ApiTimeoutError(str(exc)) from exc
+            except httpx.TransportError as exc:
+                if attempt < self.max_retries:
+                    await asyncio.sleep(retry_delay(attempt, self.retry_backoff))
+                    attempt += 1
+                    continue
+                raise ApiConnectionError(str(exc)) from exc
+
+            if response.status_code < 400:
+                return response
+
+            body = decode_json(response.text)
+            if attempt < self.max_retries and should_retry(response.status_code, body):
+                retry_after = parse_retry_after(response.headers.get("Retry-After"))
+                await asyncio.sleep(retry_delay(attempt, self.retry_backoff, retry_after))
+                attempt += 1
+                continue
+            raise error_from_response(response.status_code, response.headers, body, response.text)
+
+
+def _multipart(
+    image: ImageInput | None,
+    upload_key: str | None,
+    page: int | None,
+    extra: Mapping[str, str],
+) -> tuple[ResolvedInput | None, dict[str, Any]]:
+    """Build the ``files`` mapping for httpx; text fields are sent as parts
+    without a filename so the body is always ``multipart/form-data``."""
+
+    if (image is None) == (upload_key is None):
+        raise InvalidRequestError("Pass exactly one of `image` or `upload_key`.")
+    files: dict[str, Any] = {name: (None, value) for name, value in extra.items()}
+    if upload_key is not None:
+        files["upload_key"] = (None, upload_key)
+        if page is not None and page != 1:
+            # The object is already in storage, so the server selects the page.
+            files["page"] = (None, str(page))
+        return None, files
+    # For a file we hold, the page is cut out locally and only it is sent.
+    resolved = prepare_input(image, page)  # type: ignore[arg-type]
+    files["image"] = resolved.part
+    return resolved, files