def get_or_create(Model, **kwargs):
'''
A get_or_create method exactly like the peewee's one, but compatible with
SQLite (does not starts a transaction if SQLite is used because it does not
support nested transactions).
'''
defaults = kwargs.pop('defaults', {})
query = Model.select()
for field, value in kwargs.items():
if '__' in field:
query = query.filter(**{field: value})
else:
query = query.where(getattr(Model, field) == value)
try:
return query.get(), False
except Model.DoesNotExist:
try:
params = dict((k, v) for k, v in kwargs.items() if '__' not in k)
params.update(defaults)
if type(Model._meta.database) == SqliteDatabase:
return Model.create(**params), True
else:
with Model._meta.database.atomic():
return Model.create(**params), True
except IntegrityError as exc:
try:
return query.get(), False
except Model.DoesNotExist:
raise exc
def get_pages(
websites: tp.Iterable[tp.Iterable[str]],
session: _Session,
token_model: peewee.Model,
timeout: int = 5) -> tp.Iterable[tp.Tuple[str, tp.List[str]]]:
"""Fetches page urls from websites' main pages and stores their's tokens."""
Controller.set_page_loading_state(0, len(websites), 0)
for i, website in enumerate(websites):
website_url, page_pattern, *adt = website
try:
if not (res := session.get(website_url,
timeout=(timeout, timeout))).ok:
continue
for j, url in enumerate(
re.findall(page_pattern, res.content.decode("utf-8"))):
url = Loader._expand_url(url, website_url)
if not token_model.get_or_none(token_model.token == (
token := Loader._encode_url(url))):
token_model.create(token=token).save()
yield url, adt
Controller.set_page_loading_state(
i + 1, len(websites), j + 1)
