Skip to content

Commit 3287d05

Browse files
jcarbaughclaude
andcommitted
Add us.states.enumeration() for dynamic Enum construction
Adds a sibling to `mapping()` that builds an `Enum` keyed by `State.abbr` with member values pulled from a caller-chosen `State` attribute. Default state list matches `mapping()` (`STATES_AND_TERRITORIES`); `value_field` defaults to `name`. Closes APR-32. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 2641f2e commit 3287d05

4 files changed

Lines changed: 65 additions & 1 deletion

File tree

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
* add counties, thanks to [Ray Kiddy](https://github.com/rkiddy)
66
* add `clean_name()` method and `fallback_func` parameter to `lookup()` to provide customizable matching, thanks to [Max Filenko](https://github.com/mfilenko) and [Charlie Tonneslan](https://github.com/c-tonneslan)
77
* DC has returned to `STATES_AND_TERRITORIES`, thanks to [Kavi Gupta](https://github.com/kavigupta)
8+
* add `us.states.enumeration()` for dynamic `Enum` construction from `State` attributes
89
* add `us.__version__` and deprecate `us.version`
910
* fix `py.typed` location, thanks to [johnw-bluemark](https://github.com/johnw-bluemark)
1011
* add support for Python 3.13 and 3.14

README.md

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,34 @@ additional states argument:
234234
```
235235

236236

237+
### Enumerations
238+
239+
The `enumeration()` method dynamically constructs an `Enum` of states, keyed
240+
by state abbreviation. The value of each member is taken from the attribute
241+
named by `value_field`, which defaults to `name`.
242+
243+
```python
244+
>>> States = us.states.enumeration()
245+
>>> States.VA
246+
<States.VA: 'Virginia'>
247+
>>> States.VA.value
248+
'Virginia'
249+
250+
>>> States = us.states.enumeration('fips')
251+
>>> States.CA.value
252+
'06'
253+
```
254+
255+
Like `mapping()`, this method uses `us.STATES_AND_TERRITORIES` as the
256+
default state list, which can be overridden by passing a `states` argument:
257+
258+
```python
259+
>>> States = us.states.enumeration('fips', states=[us.states.DC, us.states.MD])
260+
>>> list(States)
261+
[<States.DC: '11'>, <States.MD: '24'>]
262+
```
263+
264+
237265
### DC should be granted statehood
238266

239267
Washington, DC does not appear in `us.STATES` or any of the

us/states.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import os
22
import re
3-
from typing import Any, Callable, Dict, Iterable, List, Optional
3+
from enum import Enum
4+
from typing import Any, Callable, Dict, Iterable, List, Optional, Type
45
from urllib.parse import urljoin
56

67
import jellyfish # type: ignore
@@ -203,6 +204,12 @@ def mapping(from_field: str, to_field: str, states: Optional[Iterable[State]] =
203204
return {getattr(s, from_field): getattr(s, to_field) for s in states}
204205

205206

207+
def enumeration(value_field: str = "name", states: Optional[Iterable[State]] = None) -> Type[Enum]:
208+
if states is None:
209+
states = STATES_AND_TERRITORIES
210+
return Enum("States", {s.abbr: getattr(s, value_field) for s in states})
211+
212+
206213
AL = State(
207214
**{
208215
"fips": "01",

us/tests/test_us.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,34 @@ def test_custom_mapping():
201201
assert "MD" in mapping
202202

203203

204+
# enumerations
205+
206+
207+
def test_enumeration():
208+
states = us.STATES[:5]
209+
enum = us.states.enumeration("name", states=states)
210+
for state in states:
211+
assert enum[state.abbr].value == state.name
212+
213+
214+
def test_enumeration_default_states():
215+
enum = us.states.enumeration("name")
216+
assert enum["VA"].value == "Virginia"
217+
assert enum["DC"].value == "District of Columbia"
218+
219+
220+
def test_enumeration_default_value_field():
221+
enum = us.states.enumeration()
222+
assert enum["VA"].value == "Virginia"
223+
224+
225+
def test_custom_enumeration():
226+
enum = us.states.enumeration("fips", states=[us.states.DC, us.states.MD])
227+
assert len(enum) == 2
228+
assert enum["DC"].value == us.states.DC.fips
229+
assert enum["MD"].value == us.states.MD.fips
230+
231+
204232
# known bugs
205233

206234

0 commit comments

Comments
 (0)