raylu 1 сар өмнө
parent
commit
5ce1fa5aca

+ 5 - 1
py/prepare.py

@@ -69,7 +69,7 @@ def add_company_ranks(totals: typing.Mapping[str, CompanyTotals]) -> typing.Mapp
 	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:
+		prod_data: typing.Mapping[str, ProdData], company_data: CompanyData) -> 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:
@@ -118,6 +118,10 @@ class ProdData(typing.TypedDict):
 	amount: float
 	volume: float
 
+class CompanyData(typing.TypedDict):
+	totals: typing.Mapping[str, CompanyTotals]
+	individual: typing.Mapping[str, typing.Mapping[str, CompanyTickerData]]
+
 class CompanyTickerData(typing.TypedDict):
 	amount: float
 	volume: float

+ 2 - 2
py/prepare_from_csv.py

@@ -19,7 +19,7 @@ def main() -> None:
 	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]]:
+		) -> 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})
@@ -40,7 +40,7 @@ def get_prod_and_company_data(data: dict[str, list[prepare.Row]], prices: typing
 			}
 			totals[row.company_id]['volume'] += volume
 
-	company_data = {'totals': prepare.add_company_ranks(totals), 'individual': dict(individual)}
+	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:

+ 2 - 2
py/prepare_from_sqlite.py

@@ -29,7 +29,7 @@ def main() -> None:
 	prepare.check_missing_tickers(prod_data)
 
 def get_prod_and_company_data(db: sqlite3.Connection, run_id: int, prices: typing.Mapping[str, float]
-		) -> tuple[typing.Mapping[str, prepare.ProdData], typing.Mapping[str, typing.Any]]:
+		) -> tuple[typing.Mapping[str, prepare.ProdData], prepare.CompanyData]:
 	cur = db.execute('''
 			SELECT ticker, entity_id, score, rank FROM leaderboard_scores
 			JOIN materials ON materials.material_pk = leaderboard_scores.material_rowid
@@ -58,7 +58,7 @@ def get_prod_and_company_data(db: sqlite3.Connection, run_id: int, prices: typin
 		amount = total_score / 30
 		prod[ticker] = {'amount': amount, 'volume': amount * prices[ticker]}
 
-	company_data = {'totals': prepare.add_company_ranks(company_totals), 'individual': dict(individual)}
+	company_data: prepare.CompanyData = {'totals': prepare.add_company_ranks(company_totals), 'individual': dict(individual)}
 	return prod, company_data
 
 if __name__ == '__main__':