新增获取省市区接口

video-gen-api/app/api/v1/pre_test_template.py

@@ -1,10 +1,13 @@
 import json
-from typing import Any, Optional
+from typing import Any, Optional, Dict
 
 from fastapi import APIRouter, Depends, HTTPException, Query, status
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy import select
 from app.models.pre_test_template import PreTestTemplate
+from app.utils.douyinApi import DouyinApi
+from app.utils.area import parse_district_data, get_area_by_level, get_cached_area_data, fetch_and_cache_area_data
+
 
 from app.dependencies import get_current_user, get_db
 from app.models.user import User
@@ -292,3 +295,64 @@ async def delete_template(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail=f"删除失败: {str(e)}",
         )
+
+
+@router.get(
+    "/getArea",
+    summary="获取行政区域信息",
+    description="获取指定级别的行政区域信息，支持一级、二级、三级区域,如果需要更新地区，执行：/api/pre-test-template/getArea?oauth_id=0019ecab9b8bc57d964&advertiser_id=1836693172153543",
+)
+async def get_template_area(
+    oauth_id: str = Query(None, description="授权ID选填，更新地区必填"),
+    advertiser_id: str = Query(default="1836693172153543", description="授权ID选填，更新地区必填"),
+    code: Optional[str] = Query("CN", description="行政区域编码，默认中国CN，选填"),
+    level: Optional[str] = Query("ONE_LEVEL", description="行政区域层级，可选值：ONE_LEVEL(获取省份)、TWO_LEVEL（市级）、THREE_LEVEL（区级）"),
+    parent_code: Optional[str] = Query(None, description="父级区域编码，获取二级时传一级编码，获取三级时传二级编码"),
+) -> Any:
+    try:
+        # 1. 先检查缓存是否存在
+        area_list = get_cached_area_data()
+        
+        # 2. 如果缓存不存在，调用接口获取数据并保存到缓存
+        if not area_list:
+            area_list = await fetch_and_cache_area_data(oauth_id, advertiser_id, code)
+        
+        # 3. 根据 level 参数过滤区域
+        if level == "ONE_LEVEL":
+            result = get_area_by_level(area_list, "ONE_LEVEL")
+        elif level == "TWO_LEVEL":
+            result = get_area_by_level(area_list, "TWO_LEVEL", parent_code)
+        elif level == "THREE_LEVEL":
+            if not parent_code:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail="获取三级区域需要提供二级区域编码(parent_code)",
+                )
+            result = get_area_by_level(area_list, "THREE_LEVEL", parent_code)
+        else:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail=f"不支持的级别: {level}",
+            )
+        
+        # 4. 转换为字典格式返回
+        result_dict = [area.to_dict() for area in result]
+        
+        return {
+            "code": 0,
+            "message": "成功",
+            "data": result_dict,
+        }
+        
+    except HTTPException:
+        raise
+    except ValueError as e:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=str(e),
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"获取区域信息失败: {str(e)}",
+        )

video-gen-api/app/utils/area.py

@@ -0,0 +1,241 @@
+import json
+import os
+from typing import Optional, List, Dict
+
+# 缓存文件路径
+CACHE_FILE_PATH = os.path.join(os.path.dirname(__file__), 'area_cache.json')
+
+
+class AreaInfo:
+    def __init__(self, code: str, name: str, level: str, geoname_id: int = None):
+        self.code = code
+        self.name = name
+        self.level = level
+        self.geoname_id = geoname_id
+        self.sub_districts: List[AreaInfo] = []
+
+    def to_dict(self):
+        return {
+            "code": self.code,
+            "name": self.name,
+            "level": self.level,
+            "geoname_id": self.geoname_id,
+            "sub_districts": [sd.to_dict() for sd in self.sub_districts] if self.sub_districts else [],
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        area = cls(
+            code=data.get("code"),
+            name=data.get("name"),
+            level=data.get("level"),
+            geoname_id=data.get("geoname_id"),
+        )
+        sub_districts = data.get("sub_districts", [])
+        for sub in sub_districts:
+            area.sub_districts.append(cls.from_dict(sub))
+        return area
+
+
+def parse_district_data(district_data: dict) -> AreaInfo:
+    """解析巨量接口返回的区域数据"""
+    area = AreaInfo(
+        code=district_data.get("code"),
+        name=district_data.get("name"),
+        level=district_data.get("level"),
+        geoname_id=district_data.get("geoname_id"),
+    )
+    
+    sub_districts = district_data.get("sub_districts")
+    if sub_districts:
+        for sub in sub_districts:
+            area.sub_districts.append(parse_district_data(sub))
+    
+    return area
+
+
+def save_area_cache(areas: List[AreaInfo]) -> None:
+    """将区域数据保存到缓存文件"""
+    data = [area.to_dict() for area in areas]
+    with open(CACHE_FILE_PATH, 'w', encoding='utf-8') as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+
+
+def load_area_cache() -> Optional[List[AreaInfo]]:
+    """从缓存文件加载区域数据"""
+    if not os.path.exists(CACHE_FILE_PATH):
+        return None
+    
+    try:
+        with open(CACHE_FILE_PATH, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+            return [AreaInfo.from_dict(item) for item in data]
+    except (json.JSONDecodeError, IOError):
+        return None
+
+
+def clear_area_cache() -> None:
+    """清除缓存文件"""
+    if os.path.exists(CACHE_FILE_PATH):
+        os.remove(CACHE_FILE_PATH)
+
+
+def filter_by_level(areas: List[AreaInfo], target_level: str) -> List[AreaInfo]:
+    """
+    根据级别过滤区域信息
+    
+    :param areas: 区域列表
+    :param target_level: ONE_LEVEL / TWO_LEVEL / THREE_LEVEL
+    :return: 指定级别的区域列表
+    """
+    result = []
+    
+    def traverse(area: AreaInfo):
+        if area.level == target_level:
+            filtered = AreaInfo(
+                code=area.code,
+                name=area.name,
+                level=area.level,
+                geoname_id=area.geoname_id,
+            )
+            result.append(filtered)
+        
+        if area.sub_districts:
+            for sub in area.sub_districts:
+                traverse(sub)
+    
+    for area in areas:
+        traverse(area)
+    
+    return result
+
+
+def get_first_level_areas(areas: List[AreaInfo]) -> List[AreaInfo]:
+    """获取所有一级区域（省/直辖市）"""
+    return filter_by_level(areas, "ONE_LEVEL")
+
+
+def get_second_level_areas(areas: List[AreaInfo], parent_code: str = None) -> List[AreaInfo]:
+    """
+    获取二级区域（市）
+    
+    :param areas: 区域列表
+    :param parent_code: 一级区域code，不传则返回所有二级区域
+    :return: 二级区域列表
+    """
+    if parent_code:
+        def find_parent_and_get_children(area: AreaInfo):
+            if area.code == parent_code:
+                return [AreaInfo(
+                    code=sub.code,
+                    name=sub.name,
+                    level=sub.level,
+                    geoname_id=sub.geoname_id,
+                ) for sub in area.sub_districts] if area.sub_districts else []
+            
+            if area.sub_districts:
+                for sub in area.sub_districts:
+                    result = find_parent_and_get_children(sub)
+                    if result:
+                        return result
+            return []
+        
+        for area in areas:
+            result = find_parent_and_get_children(area)
+            if result:
+                return result
+        return []
+    else:
+        return filter_by_level(areas, "TWO_LEVEL")
+
+
+def get_third_level_areas(areas: List[AreaInfo], parent_code: str) -> List[AreaInfo]:
+    """
+    获取三级区域（区/县）
+    
+    :param areas: 区域列表
+    :param parent_code: 二级区域code
+    :return: 三级区域列表
+    """
+    def find_parent_and_get_children(area: AreaInfo):
+        if area.code == parent_code:
+            return [AreaInfo(
+                code=sub.code,
+                name=sub.name,
+                level=sub.level,
+                geoname_id=sub.geoname_id,
+            ) for sub in area.sub_districts] if area.sub_districts else []
+        
+        if area.sub_districts:
+            for sub in area.sub_districts:
+                result = find_parent_and_get_children(sub)
+                if result:
+                    return result
+        return []
+    
+    for area in areas:
+        result = find_parent_and_get_children(area)
+        if result:
+            return result
+    return []
+
+
+def get_area_by_level(areas: List[AreaInfo], level: str, parent_code: str = None) -> List[AreaInfo]:
+    """
+    根据级别获取区域信息
+    
+    :param areas: 区域列表
+    :param level: ONE_LEVEL / TWO_LEVEL / THREE_LEVEL
+    :param parent_code: 父级区域code（TWO_LEVEL和THREE_LEVEL时可选/必填）
+    :return: 区域列表
+    """
+    level = level.upper()
+    
+    if level == "ONE_LEVEL":
+        return get_first_level_areas(areas)
+    elif level == "TWO_LEVEL":
+        return get_second_level_areas(areas, parent_code)
+    elif level == "THREE_LEVEL":
+        if not parent_code:
+            raise ValueError("获取三级区域需要提供二级区域code")
+        return get_third_level_areas(areas, parent_code)
+    else:
+        raise ValueError(f"不支持的级别: {level}")
+
+
+async def fetch_and_cache_area_data(oauth_id: str, advertiser_id: str = 1836693172153543, code: str = "CN") -> List[AreaInfo]:
+    """
+    从接口获取区域数据并缓存到文件
+    
+    :param oauth_id: 授权ID
+    :param code: 行政区域编码，默认中国CN
+    :return: 区域列表
+    """
+    from app.utils.douyinApi import DouyinApi
+    
+    params = {
+        "advertiser_id": advertiser_id,
+        "codes": json.dumps([code]),
+        "language": "ZH_CN",
+        "sub_district": "THREE_LEVEL",
+        "version": "V2_3_2"
+    }
+    
+    area_response = await DouyinApi().get_area(oauth_id=oauth_id, params=params)
+    
+    if area_response.get("code") != 0:
+        raise Exception(f"获取区域信息失败: {area_response.get('message', '未知错误')}")
+    
+    districts_data = area_response.get("data", {}).get("districts", [])
+    if not districts_data:
+        raise Exception("接口返回的区域数据为空")
+    
+    area_list = [parse_district_data(d) for d in districts_data]
+    save_area_cache(area_list)
+    
+    return area_list
+
+
+def get_cached_area_data() -> Optional[List[AreaInfo]]:
+    """获取缓存的区域数据"""
+    return load_area_cache()

video-gen-api/app/utils/douyinApi.py

@@ -53,3 +53,16 @@ class DouyinApi:
             'POST',
             options
         )
+    
+    #获取区域信息
+    async def get_area(self, oauth_id: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+        if not oauth_id:
+            raise RuntimeError('OAuth ID is not set.')
+
+        url = "https://api.oceanengine.com/open_api/2/tools/admin/info/"
+        return await self.request.request_with_token_with_context(
+            oauth_id,
+            url,
+            'GET',
+            {'params': params or {}}
+        )