AcidMap/app/models/city.py

from sqlalchemy import Column, Integer, String, Text, Index, JSON
from sqlalchemy.dialects.postgresql import JSONB
from geoalchemy2 import Geometry
from app.database import Base
from shapely.geometry import shape, MultiPolygon, Polygon
import json

class City(Base):
    """市级行政区划地理数据表"""
    __tablename__ = "cities"

    # 字段映射字典
    FIELD_MAPPING = {
        '市名': 'name',
        '市代码': 'code',
        '省名': 'province_name',
        '省代码': 'province_code'
    }

    # 反向映射字典
    REVERSE_FIELD_MAPPING = {v: k for k, v in FIELD_MAPPING.items()}

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100), nullable=False, comment="市名")
    code = Column(Integer, nullable=False, unique=True, comment="市代码")
    province_name = Column(String(100), nullable=False, comment="省名") 
    province_code = Column(Integer, nullable=False, comment="省代码")
    
    # 使用PostGIS的几何类型来存储多边形数据
    geometry = Column(Geometry('MULTIPOLYGON', srid=4326), nullable=False, comment="市级行政区划的几何数据")
    
    # 存储完整的GeoJSON数据
    geojson = Column(JSON, nullable=False, comment="完整的GeoJSON数据") 

    # 显式定义索引
    __table_args__ = (
        Index('idx_cities_name', 'name'),
        Index('idx_cities_code', 'code'),
        Index('idx_cities_province_code', 'province_code'),
        #Index('idx_cities_geometry', 'geometry', postgresql_using='gist'),
    ) 

    @classmethod
    def from_geojson_feature(cls, feature):
        """从GeoJSON Feature创建City实例
        
        Args:
            feature: GeoJSON Feature对象
            
        Returns:
            City: 新创建的City实例
        """
        properties = feature['properties']
        # 转换中文字段名为英文
        mapped_properties = {
            cls.FIELD_MAPPING.get(k, k): v 
            for k, v in properties.items()
        }
        
        # 将GeoJSON geometry转换为EWKT格式（SRID=4326）
        geometry = feature['geometry']
        geom = shape(geometry)
        # 统一为 MultiPolygon 以匹配列类型
        if isinstance(geom, Polygon):
            geom = MultiPolygon([geom])
        wkt = f"SRID=4326;{geom.wkt}"
        
        # 创建实例
        city = cls(
            **mapped_properties,
            geometry=wkt,
            geojson=feature
        )
        return city

    def to_geojson_feature(self):
        """将City实例转换为GeoJSON Feature
        
        Returns:
            dict: GeoJSON Feature对象
        """
        properties = {}
        # 转换英文字段名为中文
        for eng_field, cn_field in self.REVERSE_FIELD_MAPPING.items():
            if hasattr(self, eng_field):
                properties[cn_field] = getattr(self, eng_field)

        return {
            'type': 'Feature',
            'properties': properties,
            'geometry': self.geometry
        }