Source code for shapelib.core.cloud.writers.flatgeobuf_cloud_writer

"""Cloud writer for FlatGeobuf files with streaming optimizations."""

from __future__ import annotations

from typing import Any, Final

from shapelib.core.cloud.cloud_path import CloudPath
from shapelib.core.cloud.writers.base_cloud_writer import BaseCloudWriter
from shapelib.factories.cloud_registry import register_cloud_writer




@register_cloud_writer
class FlatGeobufCloudWriter(BaseCloudWriter):
    """
    Cloud writer for FlatGeobuf files with streaming optimizations.

    Applies cloud-optimized settings for FlatGeobuf uploads:
    - engine="pyogrio": Faster I/O engine
    - index=False: Skip pandas index (not needed for geo data)
    - use_arrow=True: Arrow serialization for faster writes
    - SPATIAL_INDEX=YES: Enable HTTP range request streaming
    - TEMPORARY_DIR: Use temp dir for sorting large datasets

    Attributes
    ----------
    EXTENSIONS : frozenset[str]
        File extensions handled by this writer.

    Examples
    --------
    >>> writer = FlatGeobufCloudWriter()
    >>> if writer.supports(cloud_path):
    ...     writer.write(gdf, cloud_path)

    Big-O
    -----
    O(n) where n is data size, plus network upload time.
    """

    EXTENSIONS: Final[frozenset[str]] = frozenset({".fgb", ".flatgeobuf"})



    def supports(self, cloud_path: CloudPath) -> bool:
        """
        Check if this writer handles the given path.

        Parameters
        ----------
        cloud_path : CloudPath
            The parsed destination path.

        Returns
        -------
        bool
            True if path is remote and extension is FlatGeobuf.
        """
        if not cloud_path.is_remote:
            return False
        return cloud_path.extension.lower() in self.EXTENSIONS


    def _build_write_kwargs(
        self,
        temp_dir: str,
        driver: str | None,
        schema: Any | None,
        index: Any | None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """
        Build kwargs with FlatGeobuf cloud optimizations.

        Applies optimized defaults that can be overridden by user kwargs.
        """
        write_kwargs: dict[str, Any] = {
            "driver": driver or "FlatGeobuf",
            "schema": schema,
            **kwargs,
        }

        # Apply optimized defaults (user can override via kwargs)
        write_kwargs.setdefault("engine", "pyogrio")
        write_kwargs.setdefault("index", False if index is None else index)
        write_kwargs.setdefault("use_arrow", True)

        # Merge layer_options with optimized defaults
        layer_options = write_kwargs.get("layer_options", {})
        layer_options.setdefault("SPATIAL_INDEX", "YES")
        layer_options.setdefault("TEMPORARY_DIR", temp_dir)
        write_kwargs["layer_options"] = layer_options

        return write_kwargs