from __future__ import annotations import collections import sys import typing import prepare def main() -> None: (month,) = sys.argv[1:] with open(f'rawData/{month}.csv', 'r', newline='') as f: data = prepare.read_data(f) with open(f'rawData/{month}-prices.json', 'r') as f: prices = prepare.get_prices(f) prod_data, company_data = get_prod_and_company_data(data, prices) prepare.write_data(month, data, prod_data, company_data) prepare.check_missing_tickers(prod_data) def get_prod_and_company_data(data: dict[str, list[prepare.Row]], prices: typing.Mapping[str, float] ) -> tuple[typing.Mapping[str, prepare.ProdData], typing.Mapping[str, typing.Any]]: prod: dict[str, prepare.ProdData] = {} individual: dict[str, dict[str, prepare.CompanyTickerData]] = collections.defaultdict(dict) totals: dict[str, prepare.CompanyTotals] = collections.defaultdict(lambda: {'volume': 0.0}) for section, rows in data.items(): if (ticker := get_production_ticker(section)) is None: continue price = prices[ticker] prod_amount = sum(row.num for row in rows) / 30 prod[ticker] = {'amount': prod_amount, 'volume': prod_amount * price} for row in rows: amount = row.num / 30 volume = amount * price individual[row.company_id][ticker] = { 'amount': amount, 'volume': volume, 'rank': row.rank, } totals[row.company_id]['volume'] += volume company_data = {'totals': prepare.add_company_ranks(totals), 'individual': dict(individual)} return prod, company_data def get_production_ticker(section: str) -> str | None: prefix = 'PRODUCTION_' suffix = '_DAYS_30' if not section.startswith(prefix) or not section.endswith(suffix): return None return section[len(prefix):-len(suffix)] if __name__ == '__main__': main()