refactor into prepare_from_csv.py and prepare.py

.gitignore

@@ -1,3 +1,4 @@
+__pycache__/
 config.toml
 node_modules/
 www/main.js

py/prepare.py

@@ -5,50 +5,8 @@ import collections
 import dataclasses
 import csv
 import json
-import sys
 import typing
 
-def main() -> None:
-	(month,) = sys.argv[1:]
-
-	with open(f'rawData/{month}.csv', 'r', newline='') as f:
-		data = read_data(f)
-
-	with open(f'rawData/{month}-prices.json', 'r') as f:
-		prices = get_prices(f)
-	prod_data, company_data = get_prod_and_company_data(data, prices)
-	with open(f'www/data/prod-data-{month}.json', 'w') as f:
-		json.dump(prod_data, f)
-	with open(f'www/data/company-data-{month}.json', 'w') as f:
-		json.dump(company_data, f)
-
-	bases_data: dict[str, dict[str, int]] = {r.company_id: {'bases': r.num, 'rank': r.rank} for r in data['BASES']}
-	ships_data: dict[str, dict[str, int]] = {r.company_id: {'ships': r.num, 'rank': r.rank} for r in data['SHIPS']}
-	for company_id in company_data['totals']:
-		bases_data.setdefault(company_id, {'bases': 1})
-		ships_data.setdefault(company_id, {'ships': 2})
-	with open(f'www/data/base-data-{month}.json', 'w') as f:
-		json.dump(bases_data, f)
-	with open(f'www/data/ship-data-{month}.json', 'w') as f:
-		json.dump(ships_data, f)
-
-	with open('www/data/universe-data.json', 'r+') as f:
-		universe_data = json.load(f)
-		universe_data[month] = {
-			'volume': sum(mat['volume'] for mat in prod_data.values()),
-			'profit': None,
-			'bases': sum(co['bases'] for co in bases_data.values()),
-			'companies': len(company_data['totals']),
-		}
-		f.seek(0)
-		json.dump(universe_data, f)
-		f.truncate()
-
-	response = httpx.get('https://api.fnar.net/material').raise_for_status()
-	tickers = frozenset(mat['Ticker'] for mat in response.json() if mat['Ticker'] != 'CMK')
-	if missing := tickers - prod_data.keys():
-		print('warning: missing production data for tickers', missing)
-
 def read_data(f: typing.TextIO) -> dict[str, list[Row]]:
 	data: dict[str, list[Row]] = collections.defaultdict(list)
 	reader = csv.reader(f)
@@ -103,44 +61,47 @@ def get_prices(f: typing.TextIO) -> typing.Mapping[str, float]:
 	prices.update(hardcoded_prices)
 	return prices
 
-def get_prod_and_company_data(data: dict[str, list[Row]], prices: typing.Mapping[str, float]
-		) -> tuple[typing.Mapping[str, ProdData], typing.Mapping[str, typing.Any]]:
-	prod: dict[str, ProdData] = {}
-	individual: dict[str, dict[str, CompanyTickerData]] = collections.defaultdict(dict)
-	totals: dict[str, 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': 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)]
-
 def add_company_ranks(totals: dict[str, CompanyTotals]) -> dict[str, CompanyTotals]:
 	ranked = sorted(totals.items(), key=lambda item: item[1]['volume'], reverse=True)
 	for rank, (company_id, company_totals) in enumerate(ranked, start=1):
 		company_totals['volumeRank'] = rank
 	return totals
 
+def write_data(month: str, data: dict[str, list[Row]],
+		prod_data: typing.Mapping[str, ProdData], company_data: typing.Mapping[str, dict]) -> None:
+	with open(f'www/data/prod-data-{month}.json', 'w') as f:
+		json.dump(prod_data, f)
+	with open(f'www/data/company-data-{month}.json', 'w') as f:
+		json.dump(company_data, f)
+
+	bases_data: dict[str, dict[str, int]] = {r.company_id: {'bases': r.num, 'rank': r.rank} for r in data['BASES']}
+	ships_data: dict[str, dict[str, int]] = {r.company_id: {'ships': r.num, 'rank': r.rank} for r in data['SHIPS']}
+	for company_id in company_data['totals']:
+		bases_data.setdefault(company_id, {'bases': 1})
+		ships_data.setdefault(company_id, {'ships': 2})
+	with open(f'www/data/base-data-{month}.json', 'w') as f:
+		json.dump(bases_data, f)
+	with open(f'www/data/ship-data-{month}.json', 'w') as f:
+		json.dump(ships_data, f)
+
+	with open('www/data/universe-data.json', 'r+') as f:
+		universe_data = json.load(f)
+		universe_data[month] = {
+			'volume': sum(mat['volume'] for mat in prod_data.values()),
+			'profit': None,
+			'bases': sum(co['bases'] for co in bases_data.values()),
+			'companies': len(company_data['totals']),
+		}
+		f.seek(0)
+		json.dump(universe_data, f)
+		f.truncate()
+
+def check_missing_tickers(prod_data: typing.Mapping[str, ProdData]) -> None:
+	response = httpx.get('https://api.fnar.net/material').raise_for_status()
+	tickers = frozenset(mat['Ticker'] for mat in response.json() if mat['Ticker'] != 'CMK')
+	if missing := tickers - prod_data.keys():
+		print('warning: missing production data for tickers', missing)
+
 @dataclasses.dataclass(frozen=True, slots=True, eq=False)
 class Row:
 	rank: int
@@ -164,6 +125,3 @@ class CompanyTickerData(typing.TypedDict):
 class CompanyTotals(typing.TypedDict, total=False):
 	volume: float
 	volumeRank: int
-
-if __name__ == '__main__':
-	main()

py/prepare_from_csv.py

@@ -0,0 +1,54 @@
+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()