prunstats/py/prepare_from_csv.py

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)

	prices = prepare.get_prices(month)
	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], prepare.CompanyData]:
	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: prepare.CompanyData = {'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()