242 lines
7.3 KiB
Python
242 lines
7.3 KiB
Python
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()
|