geospatial_data_converter/utils.py

import io
import os
import zipfile
from tempfile import TemporaryDirectory
from typing import BinaryIO
import geopandas as gpd

try:
 from kml_tricks import load_ge_data
except ImportError:
 from .kml_tricks import load_ge_data

output_format_dict = {
 "ESRI Shapefile": ("shp", "zip", "application/zip"), # must be zipped
 "OpenFileGDB": ("gdb", "zip", "application/zip"), # must be zipped
 "GeoJSON": ("geojson", "geojson", "application/geo+json"),
 "CSV": ("csv", "csv", "text/csv"),
 "KML": ("kml", "kml", "application/vnd.google-earth.kml+xml"),
}


def read_file(file: BinaryIO, *args, **kwargs) -> gpd.GeoDataFrame:
 """Read a file and return a GeoDataFrame"""
 basename, ext = os.path.splitext(os.path.basename(file.name))
 ext = ext.lower().strip(".")
 if ext == "zip":
 with TemporaryDirectory() as tmp_dir:
 tmp_file_path = os.path.join(tmp_dir, f"{basename}.{ext}")
 with open(tmp_file_path, "wb") as tmp_file:
 tmp_file.write(file.read())
 return gpd.read_file(
 f"zip://{tmp_file_path}",
 *args,
 engine="pyogrio",
 **kwargs,
 )
 elif ext in ("kml", "kmz"):
 with TemporaryDirectory() as tmp_dir:
 tmp_file_path = os.path.join(tmp_dir, f"{basename}.{ext}")
 print(file.name)
 print(basename)
 print(tmp_file_path)
 with open(tmp_file_path, "wb") as tmp_file:
 tmp_file.write(file.read())
 return load_ge_data(tmp_file_path)
 return gpd.read_file(file, *args, engine="pyogrio", **kwargs)


def zip_dir(directory: str) -> bytes:
 """Zip a directory and return the bytes"""
 zip_buffer = io.BytesIO()

 with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zipf:
 for root, dirs, files in os.walk(directory):
 for file in files:
 new_member = os.path.join(root, file)
 zipf.write(
 new_member,
 os.path.relpath(new_member, directory),
 )

 return zip_buffer.getvalue()


def convert(gdf: gpd.GeoDataFrame, output_name: str, output_format: str) -> bytes:
 """Convert a GeoDataFrame to the specified format"""
 with TemporaryDirectory() as tmpdir:
 out_path = os.path.join(tmpdir, output_name)
 if output_format == "CSV":
 gdf.to_csv(out_path)
 else:
 gdf.to_file(out_path, driver=output_format, engine="pyogrio")

 if output_format in ("ESRI Shapefile", "OpenFileGDB"):
 output_bytes = zip_dir(tmpdir)
 else:
 with open(out_path, "rb") as f:
 output_bytes = f.read()

 return output_bytes