""" geo_tools.io.writers ~~~~~~~~~~~~~~~~~~~~~ 统一的矢量数据写出接口,支持: - Shapefile (.shp) - GeoJSON (.geojson / .json) - GeoPackage (.gpkg) ← 支持追加图层 - File Geodatabase (.gdb) ← 通过 fiona OpenFileGDB 驱动 - FlatGeobuf (.fgb) - CSV(含 WKT 列) """ from __future__ import annotations from pathlib import Path from typing import Any, Literal import geopandas as gpd from app.utils.logger import get_logger logger = get_logger(__name__) # ── 主入口 ────────────────────────────────────────────────────────────────────── def write_vector( gdf: gpd.GeoDataFrame, path: str | Path, layer: str | None = None, driver: str | None = None, encoding: str = "utf-8", mode: Literal["w", "a"] = "w", **kwargs: Any, ) -> Path: """统一的矢量数据写出入口,自动识别格式。 Parameters ---------- gdf: 待写出的 GeoDataFrame。 path: 目标路径(文件或 `.gdb` 目录)。 layer: 图层名(GPKG、GDB 多图层格式时使用)。 driver: 强制指定 fiona 驱动名(通常不需要,自动推断)。 encoding: 字段编码,Shapefile 导出中文时常需 ``"gbk"``。 mode: ``"w"`` 覆盖写出,``"a"`` 追加图层(GPKG / GDB 支持)。 Returns ------- Path 实际写出的路径。 """ path = Path(path) suffix = path.suffix.lower() if suffix == ".csv": return _write_csv_vector(gdf, path) # 自动推断驱动 if driver is None: driver = _infer_driver(path) # 确保父目录存在(GDB 是目录,其父目录要存在) if suffix == ".gdb": path.parent.mkdir(parents=True, exist_ok=True) else: path.parent.mkdir(parents=True, exist_ok=True) write_kwargs: dict[str, Any] = { "driver": driver, "encoding": encoding, "mode": mode, **kwargs, } if layer is not None: write_kwargs["layer"] = layer logger.info( "写出矢量数据:%s(驱动:%s,图层:%s,模式:%s,要素数:%d)", path, driver, layer, mode, len(gdf), ) gdf.to_file(str(path), **write_kwargs) logger.info("写出完成:%s", path) return path # ── GDB 专用 ──────────────────────────────────────────────────────────────────── def write_gdb( gdf: gpd.GeoDataFrame, gdb_path: str | Path, layer: str, mode: Literal["w", "a"] = "w", encoding: str = "utf-8", **kwargs: Any, ) -> Path: """将 GeoDataFrame 写出到 Esri File Geodatabase(.gdb)中。 Parameters ---------- gdf: 待写出的 GeoDataFrame。 gdb_path: 目标 ``.gdb`` 目录路径(不存在时自动创建)。 layer: 图层名称(必填)。 mode: ``"w"`` 覆盖图层;``"a"`` 向已有 GDB 追加图层。 Notes ----- 写出 GDB 依赖 fiona 的 ``OpenFileGDB``(写)或 ``FileGDB``(需 ESRI 驱动)支持。 当前 fiona >= 1.9 的 ``OpenFileGDB`` 驱动已支持创建和写出,无需额外安装。 """ gdb_path = Path(gdb_path) if not layer: raise ValueError("写出 GDB 必须指定 layer 参数。") gdb_path.parent.mkdir(parents=True, exist_ok=True) logger.info("写出到 GDB:%s >> 图层 %r(模式:%s)", gdb_path.name, layer, mode) gdf.to_file( str(gdb_path), layer=layer, driver="OpenFileGDB", mode=mode, encoding=encoding, **kwargs, ) logger.info("GDB 写出完成:%s >> %s", gdb_path, layer) return gdb_path # ── GPKG 专用 ─────────────────────────────────────────────────────────────────── def write_gpkg( gdf: gpd.GeoDataFrame, gpkg_path: str | Path, layer: str, mode: Literal["w", "a"] = "w", **kwargs: Any, ) -> Path: """将 GeoDataFrame 写出为 GeoPackage 中的一个图层。 Parameters ---------- gpkg_path: 目标 ``.gpkg`` 文件路径(不存在时自动创建)。 layer: 图层名称(必填)。 mode: ``"w"`` 覆盖;``"a"`` 向已有 GPKG 追加图层。 """ gpkg_path = Path(gpkg_path) gpkg_path.parent.mkdir(parents=True, exist_ok=True) gdf.to_file(str(gpkg_path), layer=layer, driver="GPKG", mode=mode, **kwargs) logger.info("GPKG 写出完成:%s >> %s", gpkg_path, layer) return gpkg_path # ── CSV 写出 ──────────────────────────────────────────────────────────────────── def _write_csv_vector(gdf: gpd.GeoDataFrame, path: Path, **kwargs: Any) -> Path: """将 GeoDataFrame 写出为含 WKT 几何列的 CSV。""" path.parent.mkdir(parents=True, exist_ok=True) df = gdf.copy() df["geometry"] = df["geometry"].apply(lambda g: g.wkt if g is not None else None) df.to_csv(path, index=False, encoding="utf-8-sig", **kwargs) # utf-8-sig 兼容 Excel logger.info("CSV 写出完成:%s", path) return path def write_csv(gdf: gpd.GeoDataFrame, path: str | Path, **kwargs: Any) -> Path: """将 GeoDataFrame 写出为含 WKT 几何列的 CSV(公开接口)。""" return _write_csv_vector(gdf, Path(path), **kwargs) # ── 工具函数 ───────────────────────────────────────────────────────────────────── def _infer_driver(path: Path) -> str: """根据文件扩展名推断 fiona 驱动。""" _EXT_TO_DRIVER: dict[str, str] = { ".shp": "ESRI Shapefile", ".geojson": "GeoJSON", ".json": "GeoJSON", ".gpkg": "GPKG", ".gdb": "OpenFileGDB", ".kml": "KML", ".fgb": "FlatGeobuf", ".gml": "GML", ".dxf": "DXF", } suffix = path.suffix.lower() if path.is_dir() and suffix == ".gdb": return "OpenFileGDB" driver = _EXT_TO_DRIVER.get(suffix) if driver is None: raise ValueError( f"无法自动推断 fiona 驱动,未知扩展名:{suffix!r}。" f"请显式传入 driver=... 参数。" ) return driver