Source code for shapelib.core.handlers.zipped_shapefile_handler

from pathlib import Path
from typing import TYPE_CHECKING, Any
from zipfile import BadZipFile, ZipFile

import geopandas as gpd

from shapelib.core.handlers.geo_file_handler import GeoFileHandler
from shapelib.factories.registry import register_handler

if TYPE_CHECKING:
    from shapelib.core.geometries.base import BaseShapes


MANDATORY_SHAPEFILE_COMPONENTS = (".shp", ".shx", ".dbf")




@register_handler(file_extension=".zip")
class ZippedShapefileHandler(GeoFileHandler):
    """Handler for shapefiles stored inside a local .zip archive."""

    driver = "ESRI Shapefile"
    extension = ".zip"
    encoding_list = ["utf-8", "iso8859_15", "latin1", "cp1252", "ascii"]
    engines = ["pyogrio", "fiona"]

    def __init__(self, path: Path) -> None:
        super().__init__(path)
        self.validate_file()



    def validate_file(self) -> None:
        """
        Validate that the archive is readable and contains at least one .shp file.

        Raises
        ------
        ValueError
            If the archive is invalid or contains no shapefile.
        """
        try:
            with ZipFile(self.path) as zip_archive:
                shapefiles = [
                    member_name for member_name in zip_archive.namelist() if member_name.lower().endswith(".shp")
                ]
        except BadZipFile as error:
            raise ValueError(f"Invalid ZIP archive: {self.path.as_posix()}") from error

        if not shapefiles:
            raise ValueError(f"The ZIP archive does not contain a shapefile (.shp): {self.path.as_posix()}")


    def _find_missing_components(self) -> list[str]:
        """
        List the mandatory shapefile components absent from the archive.

        Returns
        -------
        list[str]
            Missing extensions, in ``MANDATORY_SHAPEFILE_COMPONENTS`` order.
        """
        with ZipFile(self.path) as zip_archive:
            member_names = [member_name.lower() for member_name in zip_archive.namelist()]

        return [
            extension
            for extension in MANDATORY_SHAPEFILE_COMPONENTS
            if not any(member_name.endswith(extension) for member_name in member_names)
        ]

    def _fiona_zip_uri(self) -> str:
        """
        Build the canonical zip URI required by Fiona.

        Returns
        -------
        str
            URI in zip:// format.
        """
        return f"zip://{self.path.resolve().as_posix()}"



    def load_data(self, **kwargs: Any) -> gpd.GeoDataFrame:
        """
        Load geospatial data from ZIP, trying engines and encodings.

        Parameters
        ----------
        **kwargs : Any
            Additional arguments passed to ``geopandas.read_file``.

        Returns
        -------
        gpd.GeoDataFrame
            Loaded geospatial data.

        Raises
        ------
        ValueError
            If a mandatory component is missing from the archive, or if loading
            fails across all engine and encoding combinations. The last
            underlying error is included in the message and chained.
        """
        missing_components = self._find_missing_components()
        if missing_components:
            raise ValueError(
                f"Zipped shapefile is missing mandatory component(s) {', '.join(missing_components)}: "
                f"{self.path.as_posix()}"
            )

        last_error: Exception | None = None

        for engine in self.engines:
            source: Path | str = self.path if engine == "pyogrio" else self._fiona_zip_uri()

            for encoding in self.encoding_list:
                try:
                    return gpd.read_file(source, encoding=encoding, engine=engine, **kwargs)
                except UnicodeDecodeError as error:
                    last_error = error
                    continue
                except Exception as error:  # noqa: BLE001
                    last_error = error
                    break

        message = f"Failed to read zipped shapefile: {self.path.as_posix()}"
        if last_error is None:
            raise ValueError(message)

        raise ValueError(f"{message}: {last_error}") from last_error




    def create_obj(self, **kwargs: Any) -> "BaseShapes":
        """
        Create geometry wrapper object from zipped shapefile data.

        Parameters
        ----------
        **kwargs : Any
            Additional arguments passed to ``load_data``.

        Returns
        -------
        BaseShapes
            Geometry-specific wrapper.
        """
        data = self.load_data(**kwargs)
        return self.create_geometry_instance(data=data, driver=self.driver)