prunstats/py/update_companies.py

from __future__ import annotations
import time

import itertools
import json
import sys
import tomllib
import typing

import httpx

def main() -> None:
	(month,) = sys.argv[1:]

	with open('config.toml', 'rb') as f:
		config = tomllib.load(f)
	fio_api_key = config['fio_api_key']

	with open(f'www/data/company-data-{month}.json', 'r') as f:
		company_ids = frozenset(json.load(f)['totals'])
	print(len(company_ids), 'companies in', month, 'data')

	with open('www/data/knownCompanies.json', 'r') as f:
		companies = json.load(f)

	now = int(time.time())
	cached: set[str] = set()
	for co in companies.values():
		co_last_update = co.get('LastUpdated')
		if co_last_update is not None and now - co_last_update < 24 * 60 * 60:
			cached.add(co['Username'].casefold())
	print(len(cached), 'cached companies')

	all_users: typing.Collection[str] = httpx.get('https://rest.fnar.net/user/allusers',
			headers={'Authorization': fio_api_key}).raise_for_status().json()
	try:
		for username in all_users:
			update_company_from_fio(companies, cached, now, username)
	finally:
		with open('www/data/knownCompanies.json', 'w') as f:
			json.dump(companies, f)

	unknown_companies = company_ids - companies.keys()
	print(len(unknown_companies), 'unknown companies')
	try:
		for batch in itertools.batched(unknown_companies, 10):
			company_data: typing.Collection[FIOCompany | None] = httpx.get('https://api.fnar.net/company/lookup',
					params={'company': batch}).raise_for_status().json()
			for fio_co in company_data:
				if fio_co is None:
					continue
				known_user = {'Username': fio_co['UserName'], 'LastUpdated': now}
				if (corp := fio_co.get('CorporationCode')) and corp != 'UNKN':
					known_user['Corporation'] = corp
				print(known_user)
				assert fio_co['CompanyId'] not in companies
				companies[fio_co['CompanyId']] = known_user
	finally:
		with open('www/data/knownCompanies.json', 'w') as f:
			json.dump(companies, f)

	fio_rest_users = frozenset(user.casefold() for user in all_users)
	known_usernames = frozenset(co['Username'].casefold() for co in companies.values() if co['Username'] is not None)
	users_without_fio_accounts = known_usernames - fio_rest_users
	print(len(users_without_fio_accounts), 'users without FIO REST accounts')
	try:
		for username in users_without_fio_accounts:
			update_company_from_fio(companies, cached, now, username)
	finally:
		with open('www/data/knownCompanies.json', 'w') as f:
			json.dump(companies, f)

def update_company_from_fio(companies: dict[str, dict[str, typing.Any]], cached: typing.AbstractSet[str], now: int,
		username: str) -> None:
	if username.casefold() in cached:
		return

	response = get_with_retry('https://rest.fnar.net/user/' + username)
	if response.status_code in (204, 404):
		print(username, 'not found')
		return
	fio_user = response.json()
	assert username.casefold() == fio_user['UserName'].casefold(), f'{username} != {fio_user["UserName"]}'
	known_user = {'Username': username, 'LastUpdated': now}
	if corp := fio_user.get('CorporationCode'):
		known_user['Corporation'] = corp
	if known_user == companies.get(fio_user['CompanyId']):
		print(username, 'had no change')
	else:
		print(known_user)
		companies[fio_user['CompanyId']] = known_user

def get_with_retry(url) -> httpx.Response:
	for attempt in range(10):
		try:
			response = httpx.get(url)
		except httpx.TransportError as e:
			print(e, 'retrying in', attempt)
			time.sleep(attempt)
			continue
		if response.status_code in (204, 404) or response.is_success:
			return response
		elif attempt == 9:
			response.raise_for_status()
		print(response.status_code, 'retrying in', attempt)
		time.sleep(attempt)
	raise AssertionError

class FIOCompany(typing.TypedDict):
	CompanyId: str
	UserName: str
	CorporationCode: str | None

if __name__ == '__main__':
	main()