use FIO API to map more company IDs

py/update_companies.py

@@ -1,17 +1,25 @@
 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)
 
@@ -49,6 +57,26 @@ def main() -> None:
 		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)
+
 def get_with_retry(url) -> httpx.Response:
 	for attempt in range(10):
 		try:
@@ -65,5 +93,10 @@ def get_with_retry(url) -> httpx.Response:
 		time.sleep(attempt)
 	raise AssertionError
 
+class FIOCompany(typing.TypedDict):
+	CompanyId: str
+	UserName: str
+	CorporationCode: str | None
+
 if __name__ == '__main__':
 	main()