giga-spatial 0.6.0__py3-none-any.whl

gigaspatial/core/io/data_store.py

@@ -0,0 +1,147 @@
+from abc import ABC, abstractmethod
+from typing import Any, List, Generator, Union
+
+
+class DataStore(ABC):
+    """
+    Abstract base class defining the interface for data store implementations.
+    This class serves as a parent for both local and cloud-based storage solutions.
+    """
+
+    @abstractmethod
+    def read_file(self, path: str) -> Any:
+        """
+        Read contents of a file from the data store.
+
+        Args:
+            path: Path to the file to read
+
+        Returns:
+            Contents of the file
+
+        Raises:
+            IOError: If file cannot be read
+        """
+        pass
+
+    @abstractmethod
+    def write_file(self, path: str, data: Any) -> None:
+        """
+        Write data to a file in the data store.
+
+        Args:
+            path: Path where to write the file
+            data: Data to write to the file
+
+        Raises:
+            IOError: If file cannot be written
+        """
+        pass
+
+    @abstractmethod
+    def file_exists(self, path: str) -> bool:
+        """
+        Check if a file exists in the data store.
+
+        Args:
+            path: Path to check
+
+        Returns:
+            True if file exists, False otherwise
+        """
+        pass
+
+    @abstractmethod
+    def list_files(self, path: str) -> List[str]:
+        """
+        List all files in a directory.
+
+        Args:
+            path: Directory path to list
+
+        Returns:
+            List of file paths in the directory
+        """
+        pass
+
+    @abstractmethod
+    def walk(self, top: str) -> Generator:
+        """
+        Walk through directory tree, similar to os.walk().
+
+        Args:
+            top: Starting directory for the walk
+
+        Returns:
+            Generator yielding tuples of (dirpath, dirnames, filenames)
+        """
+        pass
+
+    @abstractmethod
+    def open(self, file: str, mode: str = "r") -> Union[str, bytes]:
+        """
+        Context manager for file operations.
+
+        Args:
+            file: Path to the file
+            mode: File mode ('r', 'w', 'rb', 'wb')
+
+        Yields:
+            File-like object
+
+        Raises:
+            IOError: If file cannot be opened
+        """
+        pass
+
+    @abstractmethod
+    def is_file(self, path: str) -> bool:
+        """
+        Check if path points to a file.
+
+        Args:
+            path: Path to check
+
+        Returns:
+            True if path is a file, False otherwise
+        """
+        pass
+
+    @abstractmethod
+    def is_dir(self, path: str) -> bool:
+        """
+        Check if path points to a directory.
+
+        Args:
+            path: Path to check
+
+        Returns:
+            True if path is a directory, False otherwise
+        """
+        pass
+
+    @abstractmethod
+    def remove(self, path: str) -> None:
+        """
+        Remove a file.
+
+        Args:
+            path: Path to the file to remove
+
+        Raises:
+            IOError: If file cannot be removed
+        """
+        pass
+
+    @abstractmethod
+    def rmdir(self, dir: str) -> None:
+        """
+        Remove a directory and all its contents.
+
+        Args:
+            dir: Path to the directory to remove
+
+        Raises:
+            IOError: If directory cannot be removed
+        """
+        pass

gigaspatial/core/io/local_data_store.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+import os
+from typing import Any, List, Generator, Tuple, Union, IO
+
+from .data_store import DataStore
+
+
+class LocalDataStore(DataStore):
+    """Implementation for local filesystem storage."""
+
+    def __init__(self, base_path: Union[str, Path] = ""):
+        super().__init__()
+        self.base_path = Path(base_path).resolve()
+
+    def _resolve_path(self, path: str) -> Path:
+        """Resolve path relative to base directory."""
+        return self.base_path / path
+
+    def read_file(self, path: str) -> bytes:
+        full_path = self._resolve_path(path)
+        with open(full_path, "rb") as f:
+            return f.read()
+
+    def write_file(self, path: str, data: Union[bytes, str]) -> None:
+        full_path = self._resolve_path(path)
+        self.mkdir(str(full_path.parent), exist_ok=True)
+
+        if isinstance(data, str):
+            mode = "w"
+            encoding = "utf-8"
+        else:
+            mode = "wb"
+            encoding = None
+
+        with open(full_path, mode, encoding=encoding) as f:
+            f.write(data)
+
+    def file_exists(self, path: str) -> bool:
+        return self._resolve_path(path).is_file()
+
+    def list_files(self, path: str) -> List[str]:
+        full_path = self._resolve_path(path)
+        return [
+            str(f.relative_to(self.base_path))
+            for f in full_path.iterdir()
+            if f.is_file()
+        ]
+
+    def walk(self, top: str) -> Generator[Tuple[str, List[str], List[str]], None, None]:
+        full_path = self._resolve_path(top)
+        for root, dirs, files in os.walk(full_path):
+            rel_root = str(Path(root).relative_to(self.base_path))
+            yield rel_root, dirs, files
+
+    def list_directories(self, path: str) -> List[str]:
+        full_path = self._resolve_path(path)
+
+        if not full_path.exists():
+            return []
+
+        if not full_path.is_dir():
+            return []
+
+        return [d.name for d in full_path.iterdir() if d.is_dir()]
+
+    def open(self, path: str, mode: str = "r") -> IO:
+        full_path = self._resolve_path(path)
+        self.mkdir(str(full_path.parent), exist_ok=True)
+        return open(full_path, mode)
+
+    def is_file(self, path: str) -> bool:
+        return self._resolve_path(path).is_file()
+
+    def is_dir(self, path: str) -> bool:
+        return self._resolve_path(path).is_dir()
+
+    def remove(self, path: str) -> None:
+        full_path = self._resolve_path(path)
+        if full_path.is_file():
+            os.remove(full_path)
+
+    def rmdir(self, directory: str) -> None:
+        full_path = self._resolve_path(directory)
+        if full_path.is_dir():
+            os.rmdir(full_path)
+
+    def mkdir(self, path: str, exist_ok: bool = False) -> None:
+        full_path = self._resolve_path(path)
+        full_path.mkdir(parents=True, exist_ok=exist_ok)
+
+    def exists(self, path: str) -> bool:
+        return self._resolve_path(path).exists()

gigaspatial/core/io/readers.py

@@ -0,0 +1,265 @@
+import pandas as pd
+import geopandas as gpd
+from pathlib import Path
+import json
+import io
+import zipfile
+import gzip
+
+from .data_store import DataStore
+
+
+def read_json(data_store: DataStore, path, **kwargs):
+    with data_store.open(path, "r") as f:
+        return json.load(f, **kwargs)
+
+
+def read_kmz(file_obj, **kwargs):
+    """Helper function to read KMZ files and return a GeoDataFrame."""
+    try:
+        with zipfile.ZipFile(file_obj) as kmz:
+            # Find the KML file in the archive (usually doc.kml)
+            kml_filename = next(
+                name for name in kmz.namelist() if name.endswith(".kml")
+            )
+
+            # Read the KML content
+            kml_content = io.BytesIO(kmz.read(kml_filename))
+
+            gdf = gpd.read_file(kml_content)
+
+            # Validate the GeoDataFrame
+            if gdf.empty:
+                raise ValueError(
+                    "The KML file is empty or does not contain valid geospatial data."
+                )
+
+        return gdf
+
+    except zipfile.BadZipFile:
+        raise ValueError("The provided file is not a valid KMZ file.")
+    except StopIteration:
+        raise ValueError("No KML file found in the KMZ archive.")
+    except Exception as e:
+        raise RuntimeError(f"An error occurred: {e}")
+
+
+def read_gzipped_json_or_csv(file_path, data_store):
+    """Reads a gzipped file, attempting to parse it as JSON (lines=True) or CSV."""
+
+    with data_store.open(file_path, "rb") as f:
+        g = gzip.GzipFile(fileobj=f)
+        text = g.read().decode("utf-8")
+        try:
+            df = pd.read_json(io.StringIO(text), lines=True)
+            return df
+        except json.JSONDecodeError:
+            try:
+                df = pd.read_csv(io.StringIO(text))
+                return df
+            except pd.errors.ParserError:
+                print(f"Error: Could not parse {file_path} as JSON or CSV.")
+                return None
+
+
+def read_dataset(data_store: DataStore, path: str, compression: str = None, **kwargs):
+    """
+    Read data from various file formats stored in both local and cloud-based storage.
+
+    Parameters:
+    ----------
+    data_store : DataStore
+        Instance of DataStore for accessing data storage.
+    path : str, Path
+        Path to the file in data storage.
+    **kwargs : dict
+        Additional arguments passed to the specific reader function.
+
+    Returns:
+    -------
+    pandas.DataFrame or geopandas.GeoDataFrame
+        The data read from the file.
+
+    Raises:
+    ------
+    FileNotFoundError
+        If the file doesn't exist in blob storage.
+    ValueError
+        If the file type is unsupported or if there's an error reading the file.
+    """
+
+    # Define supported file formats and their readers
+    BINARY_FORMATS = {
+        ".shp",
+        ".zip",
+        ".parquet",
+        ".gpkg",
+        ".xlsx",
+        ".xls",
+        ".kmz",
+        ".gz",
+    }
+
+    PANDAS_READERS = {
+        ".csv": pd.read_csv,
+        ".xlsx": lambda f, **kw: pd.read_excel(f, engine="openpyxl", **kw),
+        ".xls": lambda f, **kw: pd.read_excel(f, engine="xlrd", **kw),
+        ".json": pd.read_json,
+        # ".gz": lambda f, **kw: pd.read_csv(f, compression="gzip", **kw),
+    }
+
+    GEO_READERS = {
+        ".shp": gpd.read_file,
+        ".zip": gpd.read_file,
+        ".geojson": gpd.read_file,
+        ".gpkg": gpd.read_file,
+        ".parquet": gpd.read_parquet,
+        ".kmz": read_kmz,
+    }
+
+    COMPRESSION_FORMATS = {
+        ".gz": "gzip",
+        ".bz2": "bz2",
+        ".zip": "zip",
+        ".xz": "xz",
+    }
+
+    try:
+        # Check if file exists
+        if not data_store.file_exists(path):
+            raise FileNotFoundError(f"File '{path}' not found in blob storage")
+
+        path_obj = Path(path)
+        suffixes = path_obj.suffixes
+        file_extension = suffixes[-1].lower() if suffixes else ""
+
+        if compression is None and file_extension in COMPRESSION_FORMATS:
+            compression_format = COMPRESSION_FORMATS[file_extension]
+
+            # if file has multiple extensions (e.g., .csv.gz), get the inner format
+            if len(suffixes) > 1:
+                inner_extension = suffixes[-2].lower()
+
+                if inner_extension == ".tar":
+                    raise ValueError(
+                        "Tar archives (.tar.gz) are not directly supported"
+                    )
+
+                if inner_extension in PANDAS_READERS:
+                    try:
+                        with data_store.open(path, "rb") as f:
+                            return PANDAS_READERS[inner_extension](
+                                f, compression=compression_format, **kwargs
+                            )
+                    except Exception as e:
+                        raise ValueError(f"Error reading compressed file: {str(e)}")
+                elif inner_extension in GEO_READERS:
+                    try:
+                        with data_store.open(path, "rb") as f:
+                            if compression_format == "gzip":
+                                import gzip
+
+                                decompressed_data = gzip.decompress(f.read())
+                                import io
+
+                                return GEO_READERS[inner_extension](
+                                    io.BytesIO(decompressed_data), **kwargs
+                                )
+                            else:
+                                raise ValueError(
+                                    f"Compression format {compression_format} not supported for geo data"
+                                )
+                    except Exception as e:
+                        raise ValueError(f"Error reading compressed geo file: {str(e)}")
+            else:
+                # if just .gz without clear inner type, assume csv
+                try:
+                    with data_store.open(path, "rb") as f:
+                        return pd.read_csv(f, compression=compression_format, **kwargs)
+                except Exception as e:
+                    raise ValueError(
+                        f"Error reading compressed file as CSV: {str(e)}. "
+                        f"If not a CSV, specify the format in the filename (e.g., .json.gz)"
+                    )
+
+        # Special handling for compressed files
+        if file_extension == ".zip":
+            # For zip files, we need to use binary mode
+            with data_store.open(path, "rb") as f:
+                return gpd.read_file(f)
+
+        # Determine if we need binary mode based on file type
+        mode = "rb" if file_extension in BINARY_FORMATS else "r"
+
+        # Try reading with appropriate reader
+        if file_extension in PANDAS_READERS:
+            try:
+                with data_store.open(path, mode) as f:
+                    return PANDAS_READERS[file_extension](f, **kwargs)
+            except Exception as e:
+                raise ValueError(f"Error reading file with pandas: {str(e)}")
+
+        if file_extension in GEO_READERS:
+            try:
+                with data_store.open(path, "rb") as f:
+                    return GEO_READERS[file_extension](f, **kwargs)
+            except Exception as e:
+                # For parquet files, try pandas reader if geopandas fails
+                if file_extension == ".parquet":
+                    try:
+                        with data_store.open(path, "rb") as f:
+                            return pd.read_parquet(f, **kwargs)
+                    except Exception as e2:
+                        raise ValueError(
+                            f"Failed to read parquet with both geopandas ({str(e)}) "
+                            f"and pandas ({str(e2)})"
+                        )
+                raise ValueError(f"Error reading file with geopandas: {str(e)}")
+
+        # If we get here, the file type is unsupported
+        supported_formats = sorted(set(PANDAS_READERS.keys()) | set(GEO_READERS.keys()))
+        supported_compressions = sorted(COMPRESSION_FORMATS.keys())
+        raise ValueError(
+            f"Unsupported file type: {file_extension}\n"
+            f"Supported formats: {', '.join(supported_formats)}"
+            f"Supported compressions: {', '.join(supported_compressions)}"
+        )
+
+    except Exception as e:
+        if isinstance(e, (FileNotFoundError, ValueError)):
+            raise
+        raise RuntimeError(f"Unexpected error reading dataset: {str(e)}")
+
+
+def read_datasets(data_store: DataStore, paths, **kwargs):
+    """
+    Read multiple datasets from data storage at once.
+
+    Parameters:
+    ----------
+    data_store : DataStore
+        Instance of DataStore for accessing data storage.
+    paths : list of str
+        Paths to files in data storage.
+    **kwargs : dict
+        Additional arguments passed to read_dataset.
+
+    Returns:
+    -------
+    dict
+        Dictionary mapping paths to their corresponding DataFrames/GeoDataFrames.
+    """
+    results = {}
+    errors = {}
+
+    for path in paths:
+        try:
+            results[path] = read_dataset(data_store, path, **kwargs)
+        except Exception as e:
+            errors[path] = str(e)
+
+    if errors:
+        error_msg = "\n".join(f"- {path}: {error}" for path, error in errors.items())
+        raise ValueError(f"Errors reading datasets:\n{error_msg}")
+
+    return results

gigaspatial/core/io/writers.py

@@ -0,0 +1,128 @@
+import pandas as pd
+import geopandas as gpd
+from pathlib import Path
+import json
+import io
+
+from .data_store import DataStore
+
+
+def write_json(data, data_store: DataStore, path, **kwargs):
+    with data_store.open(path, "w") as f:
+        json.dump(data, f, **kwargs)
+
+
+def write_dataset(data, data_store: DataStore, path, **kwargs):
+    """
+    Write DataFrame or GeoDataFrame to various file formats in Azure Blob Storage.
+
+    Parameters:
+    ----------
+    data : pandas.DataFrame or geopandas.GeoDataFrame
+        The data to write to blob storage.
+    data_store : DataStore
+        Instance of DataStore for accessing data storage.
+    path : str
+        Path where the file will be written in data storage.
+    **kwargs : dict
+        Additional arguments passed to the specific writer function.
+
+    Raises:
+    ------
+    ValueError
+        If the file type is unsupported or if there's an error writing the file.
+    TypeError
+        If input data is not a DataFrame or GeoDataFrame.
+    """
+
+    # Define supported file formats and their writers
+    BINARY_FORMATS = {".shp", ".zip", ".parquet", ".gpkg", ".xlsx", ".xls"}
+
+    PANDAS_WRITERS = {
+        ".csv": lambda df, buf, **kw: df.to_csv(buf, **kw),
+        ".xlsx": lambda df, buf, **kw: df.to_excel(buf, engine="openpyxl", **kw),
+        ".json": lambda df, buf, **kw: df.to_json(buf, **kw),
+        ".parquet": lambda df, buf, **kw: df.to_parquet(buf, **kw),
+    }
+
+    GEO_WRITERS = {
+        ".geojson": lambda gdf, buf, **kw: gdf.to_file(buf, driver="GeoJSON", **kw),
+        ".gpkg": lambda gdf, buf, **kw: gdf.to_file(buf, driver="GPKG", **kw),
+        ".parquet": lambda gdf, buf, **kw: gdf.to_parquet(buf, **kw),
+    }
+
+    try:
+        # Input validation
+        if not isinstance(data, (pd.DataFrame, gpd.GeoDataFrame)):
+            raise TypeError("Input data must be a pandas DataFrame or GeoDataFrame")
+
+        # Get file suffix and ensure it's lowercase
+        suffix = Path(path).suffix.lower()
+
+        # Determine if we need binary mode based on file type
+        mode = "wb" if suffix in BINARY_FORMATS else "w"
+
+        # Handle different data types and formats
+        if isinstance(data, gpd.GeoDataFrame):
+            if suffix not in GEO_WRITERS:
+                supported_formats = sorted(GEO_WRITERS.keys())
+                raise ValueError(
+                    f"Unsupported file type for GeoDataFrame: {suffix}\n"
+                    f"Supported formats: {', '.join(supported_formats)}"
+                )
+
+            try:
+                with data_store.open(path, "wb") as f:
+                    GEO_WRITERS[suffix](data, f, **kwargs)
+            except Exception as e:
+                raise ValueError(f"Error writing GeoDataFrame: {str(e)}")
+
+        else:  # pandas DataFrame
+            if suffix not in PANDAS_WRITERS:
+                supported_formats = sorted(PANDAS_WRITERS.keys())
+                raise ValueError(
+                    f"Unsupported file type for DataFrame: {suffix}\n"
+                    f"Supported formats: {', '.join(supported_formats)}"
+                )
+
+            try:
+                with data_store.open(path, mode) as f:
+                    PANDAS_WRITERS[suffix](data, f, **kwargs)
+            except Exception as e:
+                raise ValueError(f"Error writing DataFrame: {str(e)}")
+
+    except Exception as e:
+        if isinstance(e, (TypeError, ValueError)):
+            raise
+        raise RuntimeError(f"Unexpected error writing dataset: {str(e)}")
+
+
+def write_datasets(data_dict, data_store: DataStore, **kwargs):
+    """
+    Write multiple datasets to data storage at once.
+
+    Parameters:
+    ----------
+    data_dict : dict
+        Dictionary mapping paths to DataFrames/GeoDataFrames.
+    data_store : DataStore
+        Instance of DataStore for accessing data storage.
+    **kwargs : dict
+        Additional arguments passed to write_dataset.
+
+    Raises:
+    ------
+    ValueError
+        If there are any errors writing the datasets.
+    """
+    errors = {}
+
+    for path, data in data_dict.items():
+        try:
+            write_dataset(data, data_store, path, **kwargs)
+        except Exception as e:
+            errors[path] = str(e)
+
+    if errors:
+        error_msg = "\n".join(f"- {path}: {error}" for path, error in errors.items())
+        raise ValueError(f"Errors writing datasets:\n{error_msg}")