Source code for shapelib.core.geometries.operations_mixin

from collections.abc import Callable
from typing import Any, TypeVar

import geopandas as gpd

F = TypeVar("F", bound=Callable[..., Any])




def propagate_class(method: F) -> F:
    """
    Decorator to ensure that operations on a GeoDataFrame return an instance of
    the correct class and propagate metadata.

    This decorator is used to wrap methods that return a `GeoDataFrame`, ensuring
    that the resulting object is an instance of the calling class and carries over
    any metadata, such as the geometry column and CRS.

    Parameters
    ----------
    method : Callable
        The method that operates on the GeoDataFrame.

    Returns
    -------
    Callable
        A decorated method that ensures class and metadata propagation after
        GeoDataFrame operations.

    Examples
    --------
    Applying the `propagate_class` decorator to a method that modifies a GeoDataFrame:

    >>> class MyGeoDataFrame(gpd.GeoDataFrame):
    ...     @propagate_class
    ...     def custom_method(self):
    ...         # Some operation that returns a new GeoDataFrame
    ...         return gpd.GeoDataFrame()

    Notes
    -----
    This decorator is useful when extending or subclassing `geopandas.GeoDataFrame`,
    as it ensures that the result of operations like `explode` or `dissolve` remains
    an instance of the derived class, rather than reverting to a base `GeoDataFrame`.
    """

    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
        result = method(self, *args, **kwargs)
        if isinstance(result, gpd.GeoDataFrame):
            # The constructor flattens Z unless told otherwise, and __finalize__
            # copies keep_z only after it has run, so the flag goes in up front.
            extra = {"keep_z": self.keep_z} if hasattr(self, "keep_z") else {}
            return self._constructor(result, **extra).__finalize__(self)
        return result

    # Copiar metadados do método original do GeoDataFrame
    method_name = method.__name__
    if hasattr(gpd.GeoDataFrame, method_name):
        original = getattr(gpd.GeoDataFrame, method_name)
        wrapper.__doc__ = original.__doc__
        wrapper.__name__ = original.__name__
        wrapper.__qualname__ = method.__qualname__
        wrapper.__annotations__ = getattr(original, "__annotations__", {})

    return wrapper  # type: ignore[return-value]





class OperationsMixin:
    """
    A mixin class that extends common GeoDataFrame operations while ensuring class
    and metadata propagation.

    This mixin applies the `propagate_class` decorator to methods such as `explode`
    and `dissolve` to maintain the correct class type and carry over any relevant
    metadata, such as CRS and the geometry column, after performing GeoDataFrame
    operations.
    """

    @propagate_class
    def explode(self, *args: Any, **kwargs: Any) -> gpd.GeoDataFrame:
        """
        Explode multi-part geometries into single geometries.

        This method wraps `geopandas.GeoDataFrame.explode`, ensuring that the resulting
        GeoDataFrame retains the class type and metadata of the original object.

        Parameters
        ----------
        *args : tuple
            Positional arguments to pass to `geopandas.GeoDataFrame.explode`.
        **kwargs : dict
            Keyword arguments to pass to `geopandas.GeoDataFrame.explode`.

        Returns
        -------
        gpd.GeoDataFrame
            A GeoDataFrame with each multi-part geometry split into single geometries,
            maintaining the class and metadata.

        Examples
        --------
        >>> gdf = gpd.GeoDataFrame({'geometry': [MultiPolygon(...)]})
        >>> exploded_gdf = gdf.explode()

        Notes
        -----
        Use this method to break up geometries like MultiPolygons or MultiLinestrings
        into their individual components.
        """
        return super().explode(*args, **kwargs)  # type: ignore[misc]

    @propagate_class
    def overlay(self, *args: Any, **kwargs: Any) -> gpd.GeoDataFrame:
        """
        Perform a spatial overlay between two GeoDataFrames.

        This method wraps `geopandas.GeoDataFrame.overlay`, ensuring that the resulting
        GeoDataFrame retains the class type and metadata of the original object.

        Parameters
        ----------
        *args : tuple
            Positional arguments to pass to `geopandas.GeoDataFrame.overlay`.
        **kwargs : dict
            Keyword arguments to pass to `geopandas.GeoDataFrame.overlay`.

        Returns
        -------
        gpd.GeoDataFrame
            A GeoDataFrame resulting from the spatial overlay operation, maintaining
            the class and metadata.

        Examples
        --------
        Perform a spatial intersection between two GeoDataFrames:

        >>> gdf1 = gpd.GeoDataFrame({'geometry': [Polygon(...), Polygon(...)]})
        >>> gdf2 = gpd.GeoDataFrame({'geometry': [Polygon(...)]})
        >>> overlay_gdf = gdf1.overlay(gdf2, how='intersection')

        Notes
        -----
        The `overlay` function allows you to perform spatial operations between two
        GeoDataFrames, such as intersection, union, difference, or symmetric difference.
        """
        return super().overlay(*args, **kwargs)  # type: ignore[misc]