@@ -79,16 +79,51 @@ def __init__(
7979 self .subregion_name2idx = {value : key for key , value in subregion .items ()}
8080
8181
82- def unpack_poi_id (key : POI_ID , definition : _Abstract_POI_Definition ) -> tuple [int , int ]:
82+ def _resolve_poi_part (part , name2idx : dict , level_info = None ) -> int :
83+ """Resolve one half of a POI key (region or subregion) to its integer id.
84+
85+ Order: Enum -> ``.value``; numeric string -> int; other string -> ``level_info`` enum
86+ name (when ``level_info`` is a concrete enum), then ``name2idx``. When ``level_info`` is
87+ the ``Any`` wildcard and an Enum is used, a warning is emitted (the level info is not set).
88+ """
89+ if isinstance (part , Enum ):
90+ if level_info is Any :
91+ import warnings
92+
93+ warnings .warn (
94+ f"Indexing a POI with the enum { part !r} but the matching level_one_info/level_two_info is not set "
95+ f"(it is 'Any'); the integer value { part .value } is used." ,
96+ stacklevel = 4 ,
97+ )
98+ return part .value
99+ if isinstance (part , str ):
100+ try :
101+ return int (part )
102+ except ValueError :
103+ pass
104+ if level_info is not None and level_info is not Any :
105+ try :
106+ return level_info ._get_id (part , no_raise = False )
107+ except Exception :
108+ pass
109+ return name2idx [part ]
110+ return part
111+
112+
113+ def unpack_poi_id (key : POI_ID , definition : _Abstract_POI_Definition , level_one_info = None , level_two_info = None ) -> tuple [int , int ]:
83114 """Convert any supported POI key type to a ``(region, subregion)`` integer pair.
84115
85116 Accepted key forms: plain integer (packed label), ``slice(region, subregion)``,
86117 2-tuple of ints, 2-tuple of ``Abstract_lvl`` / ``Enum`` members, or mixed tuples.
87- String values are resolved via ``definition``'s name-to-index mappings.
118+ String values resolve via ``level_one_info``/``level_two_info`` enum names (when given)
119+ and otherwise via ``definition``'s name-to-index mappings, so ``poi[idx, "level2name"]``
120+ and ``poi["level1name", "level2name"]`` work in addition to ids and Enums.
88121
89122 Args:
90123 key: POI identifier in any of the supported formats.
91124 definition: Name-to-integer mapping used to resolve string labels.
125+ level_one_info: Optional region-level enum class for name/enum resolution + warnings.
126+ level_two_info: Optional subregion-level enum class for name/enum resolution + warnings.
92127
93128 Returns:
94129 ``(region, subregion)`` tuple of plain Python integers.
@@ -101,20 +136,8 @@ def unpack_poi_id(key: POI_ID, definition: _Abstract_POI_Definition) -> tuple[in
101136 subregion = key .stop
102137 else :
103138 region , subregion = key
104- if isinstance (region , str ):
105- try :
106- region = int (region )
107- except ValueError :
108- region = definition .region_name2idx [region ]
109- if isinstance (region , Enum ):
110- region = region .value
111- if isinstance (subregion , str ):
112- try :
113- subregion = int (subregion )
114- except ValueError :
115- subregion = definition .subregion_name2idx [subregion ]
116- if isinstance (subregion , Enum ):
117- subregion = subregion .value
139+ region = _resolve_poi_part (region , definition .region_name2idx , level_one_info )
140+ subregion = _resolve_poi_part (subregion , definition .subregion_name2idx , level_two_info )
118141 return region , subregion
119142
120143
@@ -731,17 +754,20 @@ def __iter__(self):
731754 """Iterate over all ``(region, subregion)`` key pairs."""
732755 return iter (self .centroids .keys ())
733756
757+ def _resolve_key (self , key : POI_ID ) -> tuple [int , int ]:
758+ """Resolve any key (ids, Enums, or ``level_one_info``/``level_two_info`` names) to ``(region, subregion)`` ints."""
759+ return unpack_poi_id (key , self .centroids .definition , self .level_one_info , self .level_two_info )
760+
734761 def __contains__ (self , key : POI_ID ) -> bool :
735- key = unpack_poi_id (key , self .centroids .definition )
736- return key in self .centroids
762+ return self ._resolve_key (key ) in self .centroids
737763
738764 def __getitem__ (self , key : POI_ID ) -> COORDINATE :
739- return tuple (self .centroids [key ])
765+ return tuple (self .centroids [self . _resolve_key ( key ) ])
740766
741767 def __setitem__ (self , key : POI_ID , value : tuple [float , float , float ] | Sequence [float ] | np .ndarray ) -> None :
742768 if len (value ) != DIMENSIONS :
743769 raise ValueError (value )
744- self .centroids [key ] = tuple (value )
770+ self .centroids [self . _resolve_key ( key ) ] = tuple (value )
745771
746772 def __len__ (self ) -> int :
747773 return self .centroids .__len__ ()
0 commit comments