bazarr/libs/pycountry/db.py
2022-11-07 13:08:27 -05:00

147 lines
4.1 KiB
Python

# vim:fileencoding=utf-8
import json
import logging
import threading
from io import open
logger = logging.getLogger("pycountry.db")
class Data:
def __init__(self, **fields):
self._fields = fields
def __getattr__(self, key):
if key not in self._fields:
raise AttributeError
return self._fields[key]
def __setattr__(self, key, value):
if key != "_fields":
self._fields[key] = value
super(Data, self).__setattr__(key, value)
def __repr__(self):
cls_name = self.__class__.__name__
fields = ", ".join("%s=%r" % i for i in sorted(self._fields.items()))
return "%s(%s)" % (cls_name, fields)
def __dir__(self):
return dir(self.__class__) + list(self._fields)
def lazy_load(f):
def load_if_needed(self, *args, **kw):
if not self._is_loaded:
with self._load_lock:
self._load()
return f(self, *args, **kw)
return load_if_needed
class Database:
data_class_base = Data
data_class_name = None
root_key = None
no_index = []
def __init__(self, filename):
self.filename = filename
self._is_loaded = False
self._load_lock = threading.Lock()
def _load(self):
if self._is_loaded:
# Help keeping the _load_if_needed code easier
# to read.
return
self.objects = []
self.index_names = set()
self.indices = {}
self.data_class = type(
self.data_class_name, (self.data_class_base,), {}
)
with open(self.filename, "r", encoding="utf-8") as f:
tree = json.load(f)
for entry in tree[self.root_key]:
obj = self.data_class(**entry)
self.objects.append(obj)
# Inject into index.
for key, value in entry.items():
# Lookups and searches are case insensitive. Normalize
# here.
value = value.lower()
if key in self.no_index:
continue
index = self.indices.setdefault(key, {})
if value in index:
logger.debug(
"%s %r already taken in index %r and will be "
"ignored. This is an error in the databases."
% (self.data_class_name, value, key)
)
index[value] = obj
self._is_loaded = True
# Public API
@lazy_load
def __iter__(self):
return iter(self.objects)
@lazy_load
def __len__(self):
return len(self.objects)
@lazy_load
def get(self, **kw):
kw.setdefault("default", None)
default = kw.pop("default")
if len(kw) != 1:
raise TypeError("Only one criteria may be given")
field, value = kw.popitem()
if not isinstance(value, str):
raise LookupError()
# Normalize for case-insensitivity
value = value.lower()
index = self.indices[field]
try:
return index[value]
except KeyError:
# Pythonic APIs implementing get() shouldn't raise KeyErrors.
# Those are a bit unexpected and they should rather support
# returning `None` by default and allow customization.
return default
@lazy_load
def lookup(self, value):
if not isinstance(value, str):
raise LookupError()
# Normalize for case-insensitivity
value = value.lower()
# Use indexes first
for key in self.indices:
try:
return self.indices[key][value]
except LookupError:
pass
# Use non-indexed values now. Avoid going through indexed values.
for candidate in self:
for k in self.no_index:
v = candidate._fields.get(k)
if v is None:
continue
if v.lower() == value:
return candidate
raise LookupError("Could not find a record for %r" % value)