#!/usr/bin/env python import itertools import re from contextlib import contextmanager from gettext import NullTranslations from typing import ( TYPE_CHECKING, Any, Callable, Dict, Iterable, List, NamedTuple, Optional, Set, Tuple, Union, cast, overload, ) from asgiref.local import Local from django.utils.encoding import force_str from django.utils.translation import override, trans_real from typing_extensions import Literal, TypedDict from django_countries.conf import settings from .base import CountriesBase if TYPE_CHECKING: from django_stubs_ext import StrPromise try: import pyuca # type: ignore collator = pyuca.Collator() # Use UCA sorting if it's available. def sort_key(item: Tuple[str, str]) -> Any: return collator.sort_key(item[1]) except ImportError: # Fallback if the UCA sorting is not available. import unicodedata # Cheap and dirty method to sort against ASCII characters only. def sort_key(item: Tuple[str, str]) -> Any: return ( unicodedata.normalize("NFKD", item[1]) .encode("ascii", "ignore") .decode("ascii") ) _translation_state = Local() class EmptyFallbackTranslator(NullTranslations): def gettext(self, message: str) -> str: if not getattr(_translation_state, "fallback", True): # Interrupt the fallback chain. return "" return super().gettext(message) @contextmanager def no_translation_fallback(): if not settings.USE_I18N: yield return # Ensure the empty fallback translator has been installed. catalog = trans_real.catalog() original_fallback = catalog._fallback if not isinstance(original_fallback, EmptyFallbackTranslator): empty_fallback_translator = EmptyFallbackTranslator() empty_fallback_translator._fallback = original_fallback catalog._fallback = empty_fallback_translator # Set the translation state to not use a fallback while inside this context. _translation_state.fallback = False try: yield finally: _translation_state.fallback = True class ComplexCountryName(TypedDict): name: "StrPromise" names: "List[StrPromise]" alpha3: str numeric: int ioc_code: str CountryName = Union["StrPromise", ComplexCountryName] CountryCode = Union[str, int, None] class AltCodes(NamedTuple): alpha3: str numeric: Optional[int] class CountryTuple(NamedTuple): code: str name: str def __repr__(self) -> str: """ Display the repr as a standard tuple for better backwards compatibility with outputting this in a template. """ return f"({self.code!r}, {self.name!r})" class Countries(CountriesBase): """ An object containing a list of ISO3166-1 countries. Iterating this object will return the countries as namedtuples (of the country ``code`` and ``name``), sorted by name. """ _countries: Dict[str, CountryName] _alt_codes: Dict[str, AltCodes] def get_option(self, option: str): """ Get a configuration option, trying the options attribute first and falling back to a Django project setting. """ value = getattr(self, option, None) if value is not None: return value return getattr(settings, f"COUNTRIES_{option.upper()}") @property def countries(self) -> Dict[str, CountryName]: """ Return the a dictionary of countries, modified by any overriding options. The result is cached so future lookups are less work intensive. """ if not hasattr(self, "_countries"): only: "Iterable[Union[str, Tuple[str, StrPromise]]]" = self.get_option( "only" ) only_choices = True if only: # Originally used ``only`` as a dict, still supported. if not isinstance(only, dict): for item in only: if isinstance(item, str): only_choices = False break self._shadowed_names: "Dict[str, List[StrPromise]]" = {} if only and only_choices: self._countries = dict(only) # type: ignore else: # Local import so that countries aren't loaded into memory # until first used. from django_countries.data import COUNTRIES countries_dict = dict(COUNTRIES) if only: self._countries = {} for item in only: if isinstance(item, str): self._countries[item] = countries_dict[item] else: key, value = item self._countries[key] = value else: self._countries = countries_dict.copy() # type: ignore if self.get_option("common_names"): for code, name in self.COMMON_NAMES.items(): if code in self._countries: self._countries[code] = name override: Dict[str, Union[CountryName, None]] = self.get_option( "override" ) if override: _countries = cast( Dict[str, Union[CountryName, None]], self._countries.copy() ) _countries.update(override) self._countries = { code: name for code, name in _countries.items() if name is not None } if self.get_option("common_names"): for code in self.COMMON_NAMES: if code in self._countries and code not in override: self._shadowed_names[code] = [countries_dict[code]] for code, names in self.OLD_NAMES.items(): if code in self._countries and code not in override: country_shadowed = self._shadowed_names.setdefault(code, []) country_shadowed.extend(names) self.countries_first = [] first: List[str] = self.get_option("first") or [] for code in first: code = self.alpha2(code) if code in self._countries: self.countries_first.append(code) return self._countries @countries.deleter def countries(self): """ Reset the countries cache in case for some crazy reason the settings or internal options change. But surely no one is crazy enough to do that, right? """ if hasattr(self, "_countries"): del self._countries if hasattr(self, "_alt_codes"): del self._alt_codes if hasattr(self, "_ioc_codes"): del self._ioc_codes if hasattr(self, "_shadowed_names"): del self._shadowed_names @property def alt_codes(self) -> Dict[str, AltCodes]: if not hasattr(self, "_alt_codes"): # Again, local import so data is not loaded unless it's needed. from django_countries.data import ALT_CODES self._alt_codes = ALT_CODES # type: ignore altered = False for code, country in self.countries.items(): if isinstance(country, dict) and ( "alpha3" in country or "numeric" in country ): if not altered: self._alt_codes = self._alt_codes.copy() altered = True alpha3, numeric = self._alt_codes.get(code, ("", None)) if "alpha3" in country: alpha3 = country["alpha3"] if "numeric" in country: numeric = country["numeric"] self._alt_codes[code] = AltCodes(alpha3, numeric) return self._alt_codes @property def ioc_codes(self) -> Dict[str, str]: if not hasattr(self, "_ioc_codes"): from django_countries.ioc_data import ISO_TO_IOC self._ioc_codes = ISO_TO_IOC altered = False for code, country in self.countries.items(): if isinstance(country, dict) and "ioc_code" in country: if not altered: self._ioc_codes = self._ioc_codes.copy() altered = True self._ioc_codes[code] = country["ioc_code"] return self._ioc_codes @property def shadowed_names(self): if not getattr(self, "_shadowed_names", False): # Getting countries populates shadowed names. self.countries return self._shadowed_names def translate_code(self, code: str, ignore_first: Optional[List[str]] = None): """ Return translated countries for a country code. """ country = self.countries[code] if isinstance(country, dict): if "names" in country: names = country["names"] else: names = [country["name"]] else: names = [country] if ignore_first and code in ignore_first: names = names[1:] for name in names: yield self.translate_pair(code, name) def translate_pair(self, code: str, name: Optional[CountryName] = None): """ Force a country to the current activated translation. :returns: ``CountryTuple(code, translated_country_name)`` namedtuple """ if name is None: name = self.countries[code] if isinstance(name, dict): if "names" in name: fallback_names: "List[StrPromise]" = name["names"][1:] name = name["names"][0] else: fallback_names = [] name = name["name"] else: fallback_names = self.shadowed_names.get(code, []) if fallback_names: with no_translation_fallback(): country_name = force_str(name) # Check if there's an older translation available if there's no # translation for the newest name. if not country_name: for fallback_name in fallback_names: fallback_name_str = force_str(fallback_name) if fallback_name_str: country_name = fallback_name_str break if not country_name: # Use the translation's fallback country name. country_name = force_str(name) else: country_name = force_str(name) return CountryTuple(code, country_name) def __iter__(self): """ Iterate through countries, sorted by name. Each country record consists of a namedtuple of the two letter ISO3166-1 country ``code`` and short ``name``. The sorting happens based on the thread's current translation. Countries that are in ``settings.COUNTRIES_FIRST`` will be displayed before any sorted countries (in the order provided), and are only repeated in the sorted list if ``settings.COUNTRIES_FIRST_REPEAT`` is ``True``. The first countries can be separated from the sorted list by the value provided in ``settings.COUNTRIES_FIRST_BREAK``. """ # Initializes countries_first, so needs to happen first. countries = self.countries # Yield countries that should be displayed first. countries_first = (self.translate_pair(code) for code in self.countries_first) if self.get_option("first_sort"): countries_first = sorted(countries_first, key=sort_key) yield from countries_first if self.countries_first: first_break = self.get_option("first_break") if first_break: yield CountryTuple("", force_str(first_break)) # Force translation before sorting. ignore_first = None if self.get_option("first_repeat") else self.countries_first countries = tuple( itertools.chain.from_iterable( self.translate_code(code, ignore_first) for code in countries ) ) # Return sorted country list. yield from sorted(countries, key=sort_key) def alpha2(self, code: CountryCode) -> str: """ Return the normalized country code when passed any type of ISO 3166-1 country code. Overridden countries objects may actually have country codes that are not two characters (for example, "GB-WLS"), so the returned length of the code is not guaranteed. If no match is found, returns an empty string. """ find: Optional[Callable] code_str = force_str(code).upper() if code_str.isdigit(): lookup_numeric = int(code_str) def find(alt_codes): return alt_codes[1] == lookup_numeric elif len(code_str) == 3: lookup_alpha3 = code_str def find(alt_codes) -> bool: return alt_codes[0] == lookup_alpha3 else: find = None if find: code_str = "" for alpha2, alt_codes in self.alt_codes.items(): if find(alt_codes): code_str = alpha2 break if code_str in self.countries: return code_str return "" def name(self, code: CountryCode) -> str: """ Return the name of a country, based on the code. If no match is found, returns an empty string. """ alpha2 = self.alpha2(code) if alpha2 not in self.countries: return "" return self.translate_pair(alpha2)[1] @overload def by_name( self, country: str, *, regex: Literal[False] = False, language: str = "en", insensitive: bool = True, ) -> str: ... @overload def by_name( self, country: str, *, regex: Literal[True], language: str = "en", insensitive: bool = True, ) -> Set[str]: ... def by_name( self, country: str, *, regex: bool = False, language: str = "en", insensitive: bool = True, ) -> Union[str, Set[str]]: """ Fetch a country's ISO3166-1 two letter country code from its name. An optional language parameter is also available. Warning: This depends on the quality of the available translations. If no match is found, returns an empty string. If ``regex`` is set to True, then rather than returning a string containing the matching country code or an empty string, a set of matching country codes is returned. If ``insensitive`` is set to False (True by default), then the search will be case sensitive. ..warning:: Be cautious about relying on this returning a country code (especially with any hard-coded string) since the ISO names of countries may change over time. """ code_list = set() if regex: re_match = re.compile(country, insensitive and re.IGNORECASE) elif insensitive: country = country.lower() with override(language): for code, check_country in self.countries.items(): if isinstance(check_country, dict): if "names" in check_country: check_names: "List[StrPromise]" = check_country["names"] else: check_names = [check_country["name"]] else: check_names = [check_country] for name in check_names: if regex: if re_match.search(str(name)): code_list.add(code) else: if insensitive: if country == name.lower(): return code else: if country == name: return code if code in self.shadowed_names: for shadowed_name in self.shadowed_names[code]: if regex: if re_match.search(str(shadowed_name)): code_list.add(code) else: if insensitive: shadowed_name = shadowed_name.lower() if country == shadowed_name: return code if regex: return code_list return "" def alpha3(self, code: CountryCode) -> str: """ Return the ISO 3166-1 three letter country code matching the provided country code. If no match is found, returns an empty string. """ alpha2 = self.alpha2(code) try: alpha3 = self.alt_codes[alpha2][0] except KeyError: alpha3 = "" return alpha3 or "" @overload def numeric( self, code: Union[str, int, None], padded: Literal[False] = False ) -> Optional[int]: ... @overload def numeric( self, code: Union[str, int, None], padded: Literal[True] ) -> Optional[str]: ... def numeric(self, code: Union[str, int, None], padded: bool = False): """ Return the ISO 3166-1 numeric country code matching the provided country code. If no match is found, returns ``None``. :param padded: Pass ``True`` to return a 0-padded three character string, otherwise an integer will be returned. """ alpha2 = self.alpha2(code) try: num = self.alt_codes[alpha2][1] except KeyError: num = None if num is None: return None if padded: return "%03d" % num return num def ioc_code(self, code: CountryCode) -> str: """ Return the International Olympic Committee three letter code matching the provided ISO 3166-1 country code. If no match is found, returns an empty string. """ alpha2 = self.alpha2(code) return self.ioc_codes.get(alpha2, "") def __len__(self): """ len() used by several third party applications to calculate the length of choices. This will solve a bug related to generating fixtures. """ count = len(self.countries) # Add first countries, and the break if necessary. count += len(self.countries_first) if self.countries_first and self.get_option("first_break"): count += 1 return count def __bool__(self): return bool(self.countries) def __contains__(self, code): """ Check to see if the countries contains the given code. """ return code in self.countries def __getitem__(self, index): """ Some applications expect to be able to access members of the field choices by index. """ try: return next(itertools.islice(self.__iter__(), index, index + 1)) except TypeError: return list( itertools.islice(self.__iter__(), index.start, index.stop, index.step) ) countries = Countries()