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.models.orm_models 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
        }