1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115 | """Marshmallow pycountry fields."""
from abc import ABC
from typing import Any
import pycountry
from loguru import logger
from pycountry.db import Data, Database
from typing_extensions import override
from comicbox.fields.fields import StringField, TrapExceptionsMeta
_ALPHA_CODES = ("alpha_2", "alpha_3", "alpha_4", "name")
class PyCountryField(StringField, ABC, metaclass=TrapExceptionsMeta):
"""A pycountry value."""
DB: Database = pycountry.countries
EMPTY_CODE = ""
def __init__(
self,
*args: Any,
serialize_name: bool = False,
allow_empty: bool = False,
**kwargs: Any,
) -> None:
"""Optionally serialize with full names."""
self._serialize_name = serialize_name
self._allow_empty = allow_empty
super().__init__(*args, **kwargs)
@staticmethod
def _clean_name(name_obj: str) -> str | None:
if not name_obj:
return None
name: str | None = StringField().deserialize(name_obj)
if not name:
return None
return name.strip()
@classmethod
def _get_pycountry(cls, tag: str, name: str) -> Data | None:
"""Get pycountry object for a country or language tag."""
try:
cleaned = cls._clean_name(name)
if not cleaned:
return None
name = cleaned
# Language lookup fails for 'en' unless alpha_2 is specified.
obj: Data | None = (
cls.DB.get(alpha_2=name) if len(name) == 2 else cls.DB.lookup(name) # noqa: PLR2004
)
except Exception as exc:
logger.warning(exc)
obj = None
if obj is None:
logger.warning(f"Couldn't find {tag} for {name}")
return obj
@classmethod
def _to_alpha_code(cls, pc_obj: Any) -> str:
code = cls.EMPTY_CODE
for attr in _ALPHA_CODES:
if code := getattr(pc_obj, attr, cls.EMPTY_CODE):
break
return code
@override
def _deserialize(
self,
value: str,
attr: str,
*args: Any,
**kwargs: Any,
) -> str:
"""Return the alpha 2 encoding."""
value = super()._deserialize(value, attr, *args, **kwargs)
code = self.EMPTY_CODE
if pc_obj := self._get_pycountry(attr, value):
code = self._to_alpha_code(pc_obj)
return code
@override
def _serialize(
self,
value: str,
attr: str,
*args: Any,
**kwargs: Any,
) -> str:
"""Return the long name."""
serialized = super()._serialize(value, attr, *args, **kwargs)
if not serialized:
return self.EMPTY_CODE
code = self.EMPTY_CODE
if pc_obj := self._get_pycountry(attr, serialized):
code = pc_obj.name if self._serialize_name else self._to_alpha_code(pc_obj)
return code
class LanguageField(PyCountryField):
"""PyCountry Language Field."""
DB: Database = pycountry.languages
class CountryField(PyCountryField):
"""PyCountry Country Field."""
DB: Database = pycountry.countries
|