prunstats/py/variety.py

from __future__ import annotations
import sys

import json

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

	with open(f'www/data/company-data-{month}.json', 'r') as f:
		companies = json.load(f)['individual']
	
	variety: dict[str, int] = {}
	for co_id, stats in companies.items():
		if len(stats) > 1:
			variety[co_id] = len(stats)

	with open('www/data/knownCompanies.json', 'r') as f:
		known_companies = json.load(f)
	with open(f'www/data/base-data-{month}.json', 'r') as f:
		base_data = json.load(f)

	for rank, (co_id, count) in enumerate(sorted(variety.items(), key=lambda item: item[1], reverse=True), start=1):
		try:
			username = known_companies[co_id]['Username']
		except KeyError:
			username = 'unknown user ' + co_id
		bases = base_data[co_id]
		base_count = bases['bases']
		base_rank = bases.get('rank')
		print(f'{rank:4}: {username:20} {count:3} {base_count:2} bases (#{base_rank})')

if __name__ == '__main__':
	main()