""" 区域数据服务（基于 city.json 省/市/区三级数据） 数据格式：省(id/name/pid) → 市(id/name/pid/sons) → 区县名称字符串数组 直辖市特征：省名 == 市名（如"北京市"省下有唯一城市"北京市"） """ import json import os import logging from collections import namedtuple from typing import Optional, Tuple, Dict, Any, List, Set log = logging.getLogger("area_service") # ---------- 返回类型 ---------- AreaMatch = namedtuple('AreaMatch', ['province_id', 'city_id', 'province', 'city']) AreaMatch.__new__.__annotations__ = { 'province_id': int, 'city_id': int, 'province': str, 'city': str, } # ---------- 常量 ---------- _PROVINCE_SUFFIXES = ['壮族自治区', '回族自治区', '维吾尔自治区', '自治区', '省', '市'] _CITY_SUFFIXES = [ # 民族自治州复合后缀（按长度降序，确保最长匹配优先） '朝鲜族自治州', '土家族苗族自治州', '藏族羌族自治州', '哈尼族彝族自治州', '布依族苗族自治州', '苗族侗族自治州', '壮族苗族自治州', '傣族景颇族自治州', '蒙古族藏族自治州', '柯尔克孜自治州', '哈萨克自治州', '蒙古自治州', '藏族自治州', '彝族自治州', '白族自治州', '傣族自治州', '傈僳族自治州', '回族自治州', '自治州', '地区', '盟', '市', ] _DISTRICT_SUFFIXES = [ '自治县', '林区', '特区', '区', '县', '市', '旗', '朝鲜族自治县', '回族彝族自治县', '彝族回族苗族自治县', '哈尼族彝族傣族自治县', ] # 虚拟城市名称别名（仅加入省级→市 Trie，不加入全局 Trie 以避免多省冲突） _VIRTUAL_CITY_ALIASES = { '省直辖县级行政区划': ['省直辖县', '省辖县', '直辖县', '辖县'], '自治区直辖县级行政区划': ['兵团城市'], } # ---------- 工具函数 ---------- def _strip_suffix(name, suffixes): """尝试去掉后缀得到简称，去掉最长的匹配""" for s in sorted(suffixes, key=len, reverse=True): if name.endswith(s) and len(name) > len(s): return name[:-len(s)] return name class _TrieNode: """前缀树节点（支持多值共存）""" __slots__ = ('children', 'values') def __init__(self) -> None: self.children: Dict[str, '_TrieNode'] = {} # 多值: [(type, id, city_id), ...] # city_id: province=0, city=c_id, district=所属city_id self.values: List[Tuple[str, int, int]] = [] class _Trie: """ 统一前缀树，省/市/区 逐字链式构造 插入模式: - insert(word, type, id, city_id): 从 root 插入 - insert_from(node, word, type, id, city_id): 从指定节点继续插入 搜索规则: - 逐字遍历 addr，沿途收集最佳匹配 - 停止位检查 city_id 唯一性：多种 city_id → 歧义 → None """ _TYPE_PRIORITY = {'district': 3, 'city': 2, 'province': 1} def __init__(self) -> None: self.root = _TrieNode() # ---------- 插入 ---------- @staticmethod def _add_value(node: _TrieNode, entry_type: str, entry_id: int, city_id: int) -> None: """在节点上添加值（仅去重，冲突由 search 停止位检测）""" for t, eid, cid in node.values: if t == entry_type and eid == entry_id: return # 完全相同，忽略 node.values.append((entry_type, entry_id, city_id)) def insert(self, word: str, entry_type: str, entry_id: int, city_id: int = 0) -> List['_TrieNode']: """从 root 插入，返回叶节点列表""" node = self.root for ch in word: if ch not in node.children: node.children[ch] = _TrieNode() node = node.children[ch] _Trie._add_value(node, entry_type, entry_id, city_id) return [node] @staticmethod def insert_from(from_nodes: List['_TrieNode'], word: str, entry_type: str, entry_id: int, city_id: int = 0) -> List['_TrieNode']: """从指定节点列表继续逐字插入，返回新叶节点列表""" leaves = [] for start in from_nodes: node = start for ch in word: if ch not in node.children: node.children[ch] = _TrieNode() node = node.children[ch] _Trie._add_value(node, entry_type, entry_id, city_id) leaves.append(node) return leaves # ---------- 搜索 ---------- def _best_value(self, node: _TrieNode) -> Optional[Tuple[str, int]]: """获取节点上最佳值（district > city > province）""" best = None best_pri = 0 for t, eid, _ in node.values: pri = self._TYPE_PRIORITY.get(t, 0) if pri > best_pri: best_pri = pri best = (t, eid) return best def search(self, text: str) -> Optional[Tuple[str, int]]: """ 逐字遍历 text，返回沿途最佳匹配 (type, id) - 记录停止节点（沿途最后有值的节点） - 停止位检查：节点上所有 city_id 是否一致 - 一致 → 返回停止节点最佳值 - 不一致（多种 city_id）→ 歧义 → None """ if not text: return None node = self.root stop_node = None for ch in text: if ch not in node.children: break # 路径断了 node = node.children[ch] if self._best_value(node) is not None: stop_node = node if stop_node is None: return None # 停止位 city_id 唯一性检查 if len(stop_node.values) > 1: city_ids = set(cid for _, _, cid in stop_node.values) if len(city_ids) > 1: return None # 多种 city_id → 歧义 return self._best_value(stop_node) class AreaService: """区域数据服务，从 city.json 加载并构建多级索引""" def __init__(self, area_data_json: str, addr_prefix_json: Optional[str] = None) -> None: area_data_json = os.path.normpath(area_data_json) log.info(f'area service loading from {area_data_json}') with open(area_data_json, 'r', encoding='utf-8') as f: data = json.load(f) # 数据容器 self._province_map: Dict[int, Dict[str, Any]] = {} # province_id -> {id, name, short_name} self._city_map: Dict[int, Dict[str, Any]] = {} # city_id -> {id, name, short_name, province_id} # 层级索引 self._addr_prefix_map: Dict[str, int] = {} # 手工前缀 -> city_id self._area_trie: _Trie = _Trie() # 统一前缀树：省/市/区 -> id self._build_maps(data) self._build_indexes() # 加载可选的地址前缀映射 if addr_prefix_json: self._load_addr_prefix_map(addr_prefix_json) log.info( f'area service loaded: {len(self._province_map)} provinces, ' f'{len(self._city_map)} cities, {len(self._districts)} districts' ) # ---------- 内部构建 ---------- def _build_maps(self, data: List[Dict[str, Any]]) -> None: """从 JSON 构建省/市/区映射""" self._districts: List[Tuple[str, int]] = [] # (区县名, city_id) for prov in data: p_id = prov['id'] p_name = prov['name'] p_short = _strip_suffix(p_name, _PROVINCE_SUFFIXES) self._province_map[p_id] = { 'id': p_id, 'name': p_name, 'short_name': p_short, } for city in prov.get('sons', []): c_id = city['id'] c_name = city['name'] c_short = _strip_suffix(c_name, _CITY_SUFFIXES) self._city_map[c_id] = { 'id': c_id, 'name': c_name, 'short_name': c_short, 'province_id': p_id, } # 收集区县 for dist_name in city.get('sons', []): self._districts.append((dist_name, c_id)) def _build_indexes(self) -> None: """ 构建统一前缀树，三种插入模式： 1. 省：从 root 插入（全称 + 简称） 2. 市：从省叶节点继续插入（层级链）+ 从 root 直接插入（全局 fallback） 3. 区：从 root 直接插入（全局 fallback） """ # 预计算：直辖市的 p_id 集合（不插入省节点，直接插入市节点） municipality_pids: Set[int] = set() for c_id, city in self._city_map.items(): p_id = city['province_id'] if city['name'] == self._province_map[p_id]['name']: municipality_pids.add(p_id) # 1. 省级（全称 + 简称） # 直辖市不插入省节点，直接插入市节点（避免省/市同名导致 trie 停在 province） prov_leaves: Dict[int, List] = {} # p_id → [叶节点] for p_id, prov in self._province_map.items(): if p_id in municipality_pids: continue # 直辖市跳过省插入 leaves = self._area_trie.insert(prov['name'], 'province', p_id, 0) if prov['short_name'] and prov['short_name'] != prov['name'] and len(prov['short_name']) > 1: leaves += self._area_trie.insert(prov['short_name'], 'province', p_id, 0) prov_leaves[p_id] = leaves # 2. 市级 city_leaves: Dict[int, List] = {} # c_id → [叶节点]（用于区县层级链） for c_id, city in self._city_map.items(): if city['name'] in ('市辖区',): continue p_id = city['province_id'] c_name = city['name'] c_short = city['short_name'] # 直辖市：不经过省节点，直接从 root 插入市 if p_id in municipality_pids: direct = self._area_trie.insert(c_name, 'city', c_id, c_id) city_leaves[c_id] = direct if c_short and c_short != c_name and len(c_short) > 1: city_leaves[c_id] += self._area_trie.insert(c_short, 'city', c_id, c_id) # 虚拟城市别名 if c_name in _VIRTUAL_CITY_ALIASES: for alias in _VIRTUAL_CITY_ALIASES[c_name]: city_leaves[c_id] += self._area_trie.insert(alias, 'city', c_id, c_id) continue # 普通城市：从省叶节点继续插入（层级链：省→市） if p_id in prov_leaves and c_name: chained = self._area_trie.insert_from(prov_leaves[p_id], c_name, 'city', c_id, c_id) city_leaves[c_id] = chained if c_short and c_short != c_name and len(c_short) > 1: city_leaves[c_id] += self._area_trie.insert_from(prov_leaves[p_id], c_short, 'city', c_id, c_id) # 虚拟城市别名也从省叶节点链式插入 if c_name in _VIRTUAL_CITY_ALIASES: for alias in _VIRTUAL_CITY_ALIASES[c_name]: city_leaves[c_id] += self._area_trie.insert_from(prov_leaves[p_id], alias, 'city', c_id, c_id) # 从 root 直接插入（全局 fallback） if c_name: direct = self._area_trie.insert(c_name, 'city', c_id, c_id) if c_id not in city_leaves: city_leaves[c_id] = direct else: city_leaves[c_id] += direct # 市简称（单字简称不插入，避免歧义：如"北林区"→"北"） if c_short and c_short != c_name and len(c_short) > 1: city_leaves[c_id] += self._area_trie.insert(c_short, 'city', c_id, c_id) # 虚拟城市别名 if c_name in _VIRTUAL_CITY_ALIASES: for alias in _VIRTUAL_CITY_ALIASES[c_name]: city_leaves[c_id] += self._area_trie.insert(alias, 'city', c_id, c_id) # 3. 区级（city_id = 所属市 c_id） # 3a. 从市叶节点继续插入（层级链：市→区） for dist_name, c_id in self._districts: if c_id in city_leaves: self._area_trie.insert_from(city_leaves[c_id], dist_name, 'district', c_id, c_id) d_short = _strip_suffix(dist_name, _DISTRICT_SUFFIXES) if d_short and d_short != dist_name and len(d_short) > 1: self._area_trie.insert_from(city_leaves[c_id], d_short, 'district', c_id, c_id) # 3b. 从 root 直接插入（全局 fallback） for dist_name, c_id in self._districts: self._area_trie.insert(dist_name, 'district', c_id, c_id) d_short = _strip_suffix(dist_name, _DISTRICT_SUFFIXES) if d_short and d_short != dist_name and len(d_short) > 1: self._area_trie.insert(d_short, 'district', c_id, c_id) # ---------- 查询接口 ---------- def _resolve(self, entry_type: str, entry_id: int) -> Optional[AreaMatch]: """将 trie 匹配结果转为 AreaMatch""" if entry_type == 'city': city = self._city_map[entry_id] prov = self._province_map.get(city['province_id']) return AreaMatch( prov['id'] if prov else 0, entry_id, prov['name'] if prov else '', city['name'] ) if entry_type == 'district': city = self._city_map.get(entry_id) if city: prov = self._province_map.get(city['province_id']) return AreaMatch( prov['id'] if prov else 0, entry_id, prov['name'] if prov else '', city['name'] ) return None def search_area(self, addr: str) -> Optional[AreaMatch]: """ 从地址中解析省市信息 addr 进前缀树，出来就是结果。 """ if not addr: return None # 1. 手工配置的前缀（最高优先级） for prefix, city_id in self._addr_prefix_map.items(): if addr.startswith(prefix): city = self._city_map.get(city_id) if city: prov = self._province_map.get(city['province_id']) return AreaMatch( prov['id'] if prov else 0, city_id, prov['name'] if prov else '', city['name'] ) return AreaMatch(0, city_id, '', '') # 2. 统一前缀树匹配 result = self._area_trie.search(addr) if not result: return None entry_type, entry_id = result # 仅匹配到省（非直辖市）→ 市未知 if entry_type == 'province': return None return self._resolve(entry_type, entry_id) def get_area_info(self, area_id: int) -> Optional[Dict[str, Any]]: """ 根据区域 id 获取信息 支持省级或市级 id。 """ area_id = int(area_id) city = self._city_map.get(area_id) if city: prov = self._province_map.get(city['province_id']) return { 'id': area_id, 'name': city['name'], 'level': 'city', 'province_id': city['province_id'], 'province_name': prov['name'] if prov else '', } prov = self._province_map.get(area_id) if prov: return {'id': area_id, 'name': prov['name'], 'level': 'province'} return None def get_province_list(self) -> List[Dict[str, Any]]: """获取所有省份列表""" return [ {'id': p['id'], 'name': p['name'], 'short_name': p['short_name']} for p in self._province_map.values() ] def get_city_list(self, province_id: Optional[int] = None) -> List[Dict[str, Any]]: """获取城市列表，可按省份过滤""" if province_id: return [ {'id': c['id'], 'name': c['name'], 'short_name': c['short_name'], 'province_id': c['province_id']} for c in self._city_map.values() if c['province_id'] == province_id ] return [ {'id': c['id'], 'name': c['name'], 'short_name': c['short_name'], 'province_id': c['province_id']} for c in self._city_map.values() ] def get_district_list(self, city_id: Optional[int] = None) -> List[str]: """获取区县列表（字符串数组），可按城市 id 过滤""" if city_id is not None: return [name for name, cid in self._districts if cid == city_id] return [name for name, _ in self._districts] def batch_get_ids_by_names(self, pairs: List[Tuple[str, str]]) -> Dict[Tuple[str, str], Tuple[int, int]]: """ 批量通过 (province_name, city_name) 查询 (province_id, city_id) Args: pairs: [(province_name, city_name), ...] Returns: {(province_name, city_name): (province_id, city_id), ...} 未找到的不会出现在结果中 """ if not pairs: return {} # 构建名称反查索引 prov_name_to_id: Dict[str, int] = { prov['name']: pid for pid, prov in self._province_map.items() } city_name_to_id: Dict[Tuple[int, str], int] = { (city['province_id'], city['name']): cid for cid, city in self._city_map.items() } result: Dict[Tuple[str, str], Tuple[int, int]] = {} for p_name, c_name in pairs: if not p_name: continue p_id = prov_name_to_id.get(p_name) if p_id is None: continue c_id = city_name_to_id.get((p_id, c_name)) if c_name else 0 result[(p_name, c_name)] = (p_id, c_id or 0) return result def _load_addr_prefix_map(self, json_path: str) -> None: """ 从 JSON 文件加载手工配置的地址前缀映射 JSON 格式: [{"prefix": "xxx", "province": "省全称", "city": "市全称"}, ...] 根据省全称和市全称查找对应的城市 id """ json_path = os.path.normpath(json_path) with open(json_path, 'r', encoding='utf-8') as f: rules = json.load(f) # 构建名称反查 map prov_name_to_id: Dict[str, int] = { prov['name']: pid for pid, prov in self._province_map.items() } city_name_to_id: Dict[Tuple[int, str], int] = { (city['province_id'], city['name']): cid for cid, city in self._city_map.items() } self._addr_prefix_map = {} for rule in rules: prefix = rule.get('prefix', '') p_name = rule.get('province', '') c_name = rule.get('city', '') p_id = prov_name_to_id.get(p_name) if not p_id: log.warning(f'addr prefix rule: province not found: {p_name}') continue c_id = city_name_to_id.get((p_id, c_name)) if not c_id: log.warning(f'addr prefix rule: city not found: {p_name} {c_name}') continue self._addr_prefix_map[prefix] = c_id log.info(f'loaded {len(self._addr_prefix_map)} addr prefix rules from {json_path}')