You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
129 lines
3.4 KiB
129 lines
3.4 KiB
6 years ago
|
# vim:fileencoding=utf-8
|
||
|
|
|
||
|
|
from io import open
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
import threading
|
||
|
|
|
||
|
|
logger = logging.getLogger('pycountry.db')
|
||
|
|
|
||
|
|
try:
|
||
|
|
unicode
|
||
|
|
except NameError:
|
||
|
|
unicode = str
|
||
|
|
|
||
|
|
|
||
|
|
class Data(object):
|
||
|
|
|
||
|
|
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(object):
|
||
|
|
|
||
|
|
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():
|
||
|
|
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):
|
||
|
|
if len(kw) != 1:
|
||
|
|
raise TypeError('Only one criteria may be given')
|
||
|
|
field, value = kw.popitem()
|
||
|
|
return self.indices[field][value]
|
||
|
|
|
||
|
|
@lazy_load
|
||
|
|
def lookup(self, value):
|
||
|
|
# try relatively quick exact matches first
|
||
|
|
if isinstance(value, (str, unicode)):
|
||
|
|
value = value.lower()
|
||
|
|
|
||
|
|
for key in self.indices:
|
||
|
|
try:
|
||
|
|
return self.indices[key][value]
|
||
|
|
except LookupError:
|
||
|
|
pass
|
||
|
|
# then try slower case-insensitive lookups
|
||
|
|
for candidate in self:
|
||
|
|
for v in candidate._fields.values():
|
||
|
|
if v is None:
|
||
|
|
continue
|
||
|
|
if v.lower() == value:
|
||
|
|
return candidate
|
||
|
|
raise LookupError('Could not find a record for %r' % value)
|