""" geo_tools.io.readers ~~~~~~~~~~~~~~~~~~~~ 统一的矢量数据读取接口，支持： - Shapefile (.shp) - GeoJSON (.geojson / .json) - GeoPackage (.gpkg) - File Geodatabase (.gdb) ← 通过 fiona OpenFileGDB / ESRI FileGDB 驱动 - KML / KMZ - FlatGeobuf (.fgb) - CSV（含 WKT 或 经纬度列） 所有函数均返回 ``geopandas.GeoDataFrame``。 """ from __future__ import annotations import os from pathlib import Path from typing import Any, Generator import fiona import geopandas as gpd from app.utils.logger import get_logger from app.utils.validators import validate_vector_path logger = get_logger(__name__) # ── 主入口 ───────────────────────────────────────────────────────────────────── def read_vector( path: str | Path, layer: str | int | None = None, crs: str | int | None = None, encoding: str = "utf-8", chunk_size: int | None = None, rows: int | None = None, **kwargs: Any, ): """统一的矢量数据读取入口，自动识别文件格式。 Parameters ---------- path: 数据路径。支持文件或目录（FileGDB ``*.gdb``）。 layer: 图层名或索引（多图层格式如 GPKG、GDB 必填；单图层可省略）。 crs: 读取后强制重投影到目标 CRS（不传则保留原始 CRS）。 encoding: 属性表编码，Shapefile 中文路径常需指定 ``"gbk"``。 chunk_size: 分块大小，默认 None（一次性读取全部数据）。 【警告】：若不设置 chunk_size，大文件可能会占用大量内存。 rows: 限制读取的行数，默认 None（读取全部数据）。 用于快速预览数据，避免读取大文件的全部内容。 **kwargs: 透传给 ``geopandas.read_file`` 的额外参数。 Returns ------- gpd.GeoDataFrame 或生成器 如果 chunk_size 为 None，返回完整的 GeoDataFrame； 如果设置了 chunk_size，返回一个生成器，每次 yield 一个 GeoDataFrame 块。 示例 ----- # 全量读取（老方法） gdf = read_vector("data.shp") # 分块读取（新方法） for chunk in read_vector("large_data.shp", chunk_size=10000): # 处理每个数据块 print(f"处理了 {len(chunk)} 条数据") # 在这里做你的操作，比如计算、过滤等 # 只读取前 5 行数据（预览模式） gdf_preview = read_vector("large_data.shp", rows=5) print(gdf_preview.head()) """ path, layer = _split_gdb_layer(path) path = validate_vector_path(path) suffix = path.suffix.lower() logger.info("读取矢量数据：%s（格式：%s，图层：%s）", path, suffix or "目录", layer) if suffix == ".csv": # CSV 文件暂时不支持分块读取 if chunk_size is not None: logger.warning("CSV 文件暂不支持分块读取，将一次性读取全部数据") return _read_csv_vector(path, crs=crs, **kwargs) # fiona / geopandas 通用读取 read_kwargs: dict[str, Any] = {"encoding": encoding, **kwargs} if layer is not None: read_kwargs["layer"] = layer # 分块读取模式 if chunk_size is not None: def _chunk_generator(): logger.info("启用分块读取模式，每块 %d 条数据", chunk_size) try: # 使用 fiona 打开文件 with fiona.open(str(path), **read_kwargs) as src: # 获取坐标系信息 crs_info = src.crs # 分块读取 features = [] for i, feature in enumerate(src): # 检查是否达到行数限制 if rows is not None and i >= rows: break features.append(feature) if (i + 1) % chunk_size == 0: # 创建 GeoDataFrame 并设置 CRS gdf = gpd.GeoDataFrame.from_features(features, crs=crs_info) # 重投影 if crs is not None: gdf = gdf.to_crs(crs) # type: ignore logger.debug("读取并处理第 %d 块数据，共 %d 条", (i + 1) // chunk_size, len(gdf)) yield gdf features = [] # 处理最后一块 if features: gdf = gpd.GeoDataFrame.from_features(features, crs=crs_info) if crs is not None: gdf = gdf.to_crs(crs) # type: ignore logger.debug("读取并处理最后一块数据，共 %d 条", len(gdf)) yield gdf except Exception as exc: raise RuntimeError(f"无法分块读取矢量数据：{exc}") from None return _chunk_generator() else: # 一次性读取模式 try: # 添加 rows 参数到读取参数中 if rows is not None: read_kwargs["rows"] = rows logger.info("限制读取行数：%d", rows) gdf = gpd.read_file(str(path), **read_kwargs) except Exception as exc: raise RuntimeError(f"无法读取矢量数据：{exc}") from None if crs is not None: logger.debug("重投影到 %s", crs) gdf = gdf.to_crs(crs) # type: ignore logger.info("读取完成：共 %d 条要素，CRS=%s", len(gdf), gdf.crs) return gdf # type: ignore # ── GDB 专用 ─────────────────────────────────────────────────────────────────── def read_gdb( gdb_path: str | Path, layer: str | int | None = None, crs: str | int | None = None, encoding: str = "utf-8", **kwargs: Any, ) -> gpd.GeoDataFrame: """读取 Esri File Geodatabase（.gdb）中的图层。 Parameters ---------- gdb_path: ``.gdb`` 目录路径。 layer: 图层名称或索引。若不指定且 GDB 仅有一个图层，则自动选取第一层； 多图层时必须指定。 crs: 读取后目标 CRS，``None`` 则保留原始坐标系。 encoding: 属性表字段编码。 """ gdb_path = Path(gdb_path) if not gdb_path.exists(): raise FileNotFoundError(f"GDB 路径不存在：{gdb_path}") if gdb_path.suffix.lower() != ".gdb": raise ValueError(f"期望 .gdb 目录，收到：{gdb_path.suffix!r}") available_layers = list_gdb_layers(gdb_path) logger.debug("GDB 可用图层：%s", available_layers) if layer is None: if not available_layers: raise ValueError(f"GDB 中没有可用图层：{gdb_path}") layer = available_layers[0] if len(available_layers) > 1: logger.warning( "GDB 包含多个图层 %s，默认读取第一层 %r。请显式传入 layer=... 以指定图层。", available_layers, layer, ) logger.info("读取 GDB 图层：%s >> %s", gdb_path.name, layer) gdf = gpd.read_file(str(gdb_path), layer=layer, encoding=encoding, **kwargs) if crs is not None: gdf = gdf.to_crs(crs) # type: ignore logger.info("GDB 读取完成：%d 条要素，CRS=%s", len(gdf), gdf.crs) return gdf # type: ignore def list_gdb_layers(gdb_path: str | Path) -> list[str]: """列出 FileGDB 中所有图层名称。 Parameters ---------- gdb_path: ``.gdb`` 目录路径。 Returns ------- list[str] 图层名称列表。 """ gdb_path = Path(gdb_path) try: return fiona.listlayers(str(gdb_path)) except Exception as exc: raise RuntimeError( f"无法列出 GDB 图层：{gdb_path}。\n" "请确认 fiona 已安装 OpenFileGDB 驱动（通常随 conda/wheels 自带）。\n" f"原始错误：{exc}" ) from exc def _split_gdb_layer(path: str | Path) -> tuple[Path, str | None]: """从完整路径中分离 GDB 数据库路径和图层名。 Parameters ---------- path: 完整路径，可以是字符串或 Path 对象。 Returns ------- tuple[Path, str | None] (gdb_path, layer_name)，其中 gdb_path 是 GDB 目录路径，layer_name 是图层名，若没有图层名则为 None。 """ path_obj = Path(path) str_path = str(path_obj) # 查找 .gdb 的位置 gdb_pos = str_path.find('.gdb') if gdb_pos == -1: # 如果没有 .gdb，整个路径作为 GDB 路径，没有图层 return path_obj, None # 提取 GDB 路径（包含 .gdb） gdb_path = str_path[:gdb_pos + 4] # 提取图层名（.gdb 之后的部分） layer_part = str_path[gdb_pos + 4:] # 去除开头的路径分隔符 layer_name = layer_part.lstrip(os.sep).lstrip('/').lstrip('\\') # 如果没有图层名，返回 None if not layer_name: layer_name = None return Path(gdb_path), layer_name # ── GPKG 专用 ────────────────────────────────────────────────────────────────── def read_gpkg( gpkg_path: str | Path, layer: str | int | None = None, crs: str | int | None = None, **kwargs: Any, ) -> gpd.GeoDataFrame: """读取 GeoPackage (.gpkg) 文件。 Parameters ---------- gpkg_path: ``.gpkg`` 文件路径。 layer: 图层名或索引；多图层时必须指定。 """ gpkg_path = Path(gpkg_path) if not gpkg_path.exists(): raise FileNotFoundError(f"GPKG 文件不存在：{gpkg_path}") available = fiona.listlayers(str(gpkg_path)) if layer is None: if not available: raise ValueError(f"GPKG 中没有可用图层：{gpkg_path}") layer = available[0] if len(available) > 1: logger.warning( "GPKG 包含多个图层 %s，默认读取第一层 %r。", available, layer ) gdf = gpd.read_file(str(gpkg_path), layer=layer, **kwargs) if crs is not None: gdf = gdf.to_crs(crs) # type: ignore return gdf # type: ignore def list_gpkg_layers(gpkg_path: str | Path) -> list[str]: """列出 GeoPackage 中所有图层名称。 Parameters ---------- gpkg_path: GeoPackage 文件路径，可以是字符串或 Path 对象。 Returns ------- list[str] 图层名称列表。 """ return fiona.listlayers(str(gpkg_path)) # ── CSV 矢量读取 ──────────────────────────────────────────────────────────────── def _read_csv_vector( path: Path, lon_col: str = "longitude", lat_col: str = "latitude", wkt_col: str | None = None, crs: str | int | None = None, **kwargs: Any, ) -> gpd.GeoDataFrame: """从 CSV 读取空间数据，支持 WKT 列或经纬度列。 Parameters ---------- path: CSV 文件路径。 lon_col: 经度列名（WKT 模式时忽略）。 lat_col: 纬度列名（WKT 模式时忽略）。 wkt_col: WKT 几何列名；若指定则优先使用。 """ import pandas as pd from shapely import wkt as shapely_wkt df = pd.read_csv(path, **kwargs) if wkt_col and wkt_col in df.columns: geometry = df[wkt_col].apply(shapely_wkt.loads) elif lon_col in df.columns and lat_col in df.columns: from shapely.geometry import Point geometry = [Point(lon, lat) for lon, lat in zip(df[lon_col], df[lat_col])] else: raise ValueError( f"CSV 中未找到 WKT 列 {wkt_col!r} 或经纬度列 ({lon_col!r}, {lat_col!r})。" ) gdf = gpd.GeoDataFrame(df, geometry=geometry, crs=crs or "EPSG:4326") return gdf def read_csv_points( path: str | Path, lon_col: str = "longitude", lat_col: str = "latitude", crs: str | int = "EPSG:4326", **kwargs: Any, ) -> gpd.GeoDataFrame: """从含经纬度列的 CSV 文件创建点 GeoDataFrame（公开接口）。""" path = Path(path) return _read_csv_vector(path, lon_col=lon_col, lat_col=lat_col, crs=crs, **kwargs)