1
0
personal-data/personal_data/main.py

273 lines
7.5 KiB
Python
Raw Permalink Normal View History

2024-04-16 22:38:57 +00:00
import csv
import datetime
2024-05-15 22:47:42 +00:00
import decimal
2024-04-16 22:38:57 +00:00
import io
import logging
2024-05-21 21:37:59 +00:00
from collections.abc import Iterable, Mapping, Sequence
2024-05-15 22:29:06 +00:00
from decimal import Decimal
2024-04-16 22:38:57 +00:00
import requests
import requests_cache
from frozendict import frozendict
2024-05-09 15:58:39 +00:00
logger = logging.getLogger(__name__)
2024-04-16 22:38:57 +00:00
try:
import cfscrape
except ImportError:
cfscrape = None
2024-05-09 15:58:39 +00:00
logger.warning('cfscrape not installed: Certain fetchers might not work')
try:
import browsercookie
except ImportError:
logger.warning('browsercookie not installed: Certain fetchers might not work')
browsercookie = None
2024-04-16 22:38:57 +00:00
import personal_data.data
2024-05-17 22:33:47 +00:00
import personal_data.fetchers
2024-04-16 22:38:57 +00:00
from . import mailgun
CSV_DIALECT = 'one_true_dialect'
csv.register_dialect(CSV_DIALECT, lineterminator='\n', skipinitialspace=True)
logging.basicConfig()
logger.setLevel('INFO')
def try_value(fn, s: str) -> object:
try:
return fn(s)
2024-05-15 22:29:06 +00:00
except (ValueError, decimal.InvalidOperation):
2024-04-16 22:38:57 +00:00
return None
def to_value(s: str) -> object:
s = s.strip()
if len(s) == 0:
return None
2024-05-15 22:29:06 +00:00
if (v := try_value(Decimal, s)) is not None:
2024-04-16 22:38:57 +00:00
return v
if v := try_value(datetime.date.fromisoformat, s):
return v
if v := try_value(datetime.datetime.fromisoformat, s):
return v
if s.lower() == 'false':
return False
if s.lower() == 'true':
return True
if s.lower() == 'none':
return None
return s
2024-05-21 21:37:59 +00:00
def equals_without_fields(
a: Mapping[str, object],
b: Mapping[str, object],
fields: Iterable[str] = frozenset(),
) -> bool:
a = dict(a)
b = dict(b)
for f in fields:
del a[f], b[f]
return frozendict(a) == frozendict(b)
2024-05-18 19:35:58 +00:00
def deduplicate_dicts(
dicts: Sequence[dict],
2024-04-16 22:38:57 +00:00
deduplicate_mode: personal_data.data.DeduplicateMode,
deduplicate_ignore_columns: list[str],
2024-05-18 19:35:58 +00:00
) -> tuple[Sequence[dict], list[str]]:
2024-04-16 22:38:57 +00:00
fieldnames = []
for d in dicts:
for k in d.keys():
if k not in fieldnames:
fieldnames.append(k)
del k
del d
if deduplicate_mode == personal_data.data.DeduplicateMode.ONLY_LATEST:
2024-04-23 20:58:25 +00:00
while len(dicts) >= 2 and equals_without_fields(
dicts[-1],
dicts[-2],
deduplicate_ignore_columns,
):
2024-04-16 22:38:57 +00:00
del dicts[-1]
2024-05-21 21:37:59 +00:00
elif deduplicate_mode == personal_data.data.DeduplicateMode.BY_ALL_COLUMNS:
to_remove = set()
for idx1, first in enumerate(dicts):
for second in dicts[idx1 + 1 :]:
if equals_without_fields(first, second, deduplicate_ignore_columns):
to_remove.add(second)
dicts = set(dicts) - to_remove
2024-04-16 22:38:57 +00:00
elif deduplicate_mode != personal_data.data.DeduplicateMode.NONE:
dicts = set(dicts)
dicts = sorted(dicts, key=lambda d: tuple(str(d.get(fn, '')) for fn in fieldnames))
2024-05-18 19:35:58 +00:00
return dicts, fieldnames
2024-05-18 19:52:22 +00:00
def normalize_dict(d: dict) -> frozendict:
return frozendict(
{k: to_value(str(v)) for k, v in d.items() if to_value(str(v)) is not None},
)
2024-05-18 19:35:58 +00:00
def extend_csv_file(
filename: str,
new_dicts: list[dict],
deduplicate_mode: personal_data.data.DeduplicateMode,
deduplicate_ignore_columns: list[str],
) -> dict:
dicts = []
try:
with open(filename) as csvfile:
reader = csv.DictReader(csvfile, dialect=CSV_DIALECT)
for row in reader:
for k in list(row.keys()):
orig = row[k]
row[k] = to_value(orig)
if row[k] is None:
del row[k]
del k, orig
dicts.append(frozendict(row))
del row
del csvfile
except FileNotFoundError as e:
logger.info('Creating file: %s', filename)
original_num_dicts = len(dicts)
2024-05-18 19:52:22 +00:00
dicts += [normalize_dict(d) for d in new_dicts]
2024-05-18 19:35:58 +00:00
del new_dicts
dicts, fieldnames = deduplicate_dicts(
dicts,
deduplicate_mode,
deduplicate_ignore_columns,
)
2024-04-16 22:38:57 +00:00
csvfile_in_memory = io.StringIO()
writer = csv.DictWriter(
csvfile_in_memory,
fieldnames=fieldnames,
dialect=CSV_DIALECT,
)
writer.writeheader()
for d in dicts:
writer.writerow(d)
output_csv = csvfile_in_memory.getvalue()
del writer, csvfile_in_memory
with open(filename, 'w') as csvfile:
csvfile.write(output_csv)
del csvfile
2024-04-16 22:45:15 +00:00
logger.info(
2024-04-16 22:38:57 +00:00
'Extended CSV "%s" from %d to %d lines',
filename,
original_num_dicts,
len(dicts),
)
return {
2024-04-23 20:58:25 +00:00
'extended': original_num_dicts != len(dicts),
'input_lines': original_num_dicts,
'output_lines': len(dicts),
'dicts': dicts,
2024-04-16 22:38:57 +00:00
}
STANDARD_HEADERS = {
'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64; rv:122.0) Gecko/20100101 Firefox/122.0',
# "Accept": "application/json, text/plain, */*",
'Accept-Language': 'en-US,en;q=0.5',
'Accept-Encoding': 'gzip, deflate, br',
}
2024-04-17 22:02:13 +00:00
if cfscrape:
2024-04-23 20:58:25 +00:00
2024-04-17 22:02:13 +00:00
class CachedCfScrape(requests_cache.CacheMixin, cfscrape.CloudflareScraper):
pass
2024-04-16 22:38:57 +00:00
def get_session(cookiejar, *, with_cfscrape: bool) -> requests.Session:
assert isinstance(with_cfscrape, bool)
2024-04-17 22:02:13 +00:00
session_class = requests_cache.CachedSession
if cfscrape:
session_class = CachedCfScrape
session = session_class('web_cache', cookies=cookiejar)
2024-04-16 22:38:57 +00:00
for cookie in cookiejar:
session.cookies.set_cookie(cookie)
return session
2024-04-23 20:58:25 +00:00
def send_notification(
session: requests.Session,
scraper_name: str,
latest_dict: frozendict,
):
2024-04-16 22:45:15 +00:00
body = ['A new update has occured for ', scraper_name, '\n']
for k, v in latest_dict.items():
body.append(f'{k}: {v}\n')
mailgun.send_email(session, f'Updated {scraper_name}', ''.join(body))
2024-04-16 22:38:57 +00:00
2024-05-09 14:59:56 +00:00
2024-04-28 21:45:47 +00:00
def available_scrapers() -> list[type[personal_data.data.Scraper]]:
return personal_data.data.Scraper.__subclasses__()
2024-04-23 20:58:25 +00:00
2024-05-09 14:59:56 +00:00
2024-04-28 21:45:47 +00:00
def available_scraper_names() -> list[str]:
return [scraper_cls.__name__ for scraper_cls in available_scrapers()]
2024-05-09 14:59:56 +00:00
def main(
scraper_filter: frozenset[str],
*,
use_cookiejar: bool,
send_email_notification: bool = False,
) -> None:
2024-04-17 22:13:56 +00:00
if use_cookiejar:
cookiejar = browsercookie.firefox()
logger.info('Got cookiejar from firefox: %s cookies', len(cookiejar))
else:
cookiejar = []
logger.warning('No cookiejar is used')
2024-04-16 22:38:57 +00:00
2024-05-15 22:05:30 +00:00
if not send_email_notification:
logger.info('Email not enabled: Notifications will not be sent!')
2024-04-28 21:45:47 +00:00
for scraper_cls in available_scrapers():
2024-04-16 22:38:57 +00:00
session = get_session(cookiejar, with_cfscrape=scraper_cls.requires_cfscrape())
scraper = scraper_cls(session)
if scraper_cls.__name__ not in scraper_filter:
continue
2024-04-16 22:45:15 +00:00
logger.info(
2024-04-16 22:38:57 +00:00
'Running %s, appending to "%s"',
scraper_cls.__name__,
scraper.dataset_name,
)
2024-04-28 21:45:47 +00:00
result_rows = []
2024-04-16 22:38:57 +00:00
try:
for result in scraper.scrape():
result_rows.append(result)
del result
except requests.exceptions.HTTPError:
logger.exception('Failed in running %s', scraper_cls.__name__)
continue
status = extend_csv_file(
2024-05-18 19:35:58 +00:00
f'output/{scraper.dataset_name}.csv',
2024-04-16 22:38:57 +00:00
result_rows,
deduplicate_mode=scraper.deduplicate_mode,
deduplicate_ignore_columns=scraper.deduplicate_ignore_columns,
)
2024-04-16 22:45:15 +00:00
logger.info('Scraper done: %s', scraper.dataset_name)
2024-04-16 22:38:57 +00:00
2024-05-15 22:05:30 +00:00
if status['extended'] and send_email_notification:
send_notification(session, scraper_cls.__name__, status['dicts'][-1])
2024-04-16 22:38:57 +00:00
del scraper, session