from __future__ import annotations import httpx import collections import dataclasses import csv import json import sqlite3 import sys import typing import prepare def main() -> None: month, sqlite_path, run_id_s = sys.argv[1:] run_id = int(run_id_s) db = sqlite3.connect(sqlite_path) db.row_factory = sqlite3.Row cur = db.execute('SELECT started_at, finished_at, type, range, materials_failed FROM runs WHERE id = ?', (run_id,)) (run,) = cur.fetchall() print(', '.join(f'{k}: {run[k]}' for k in run.keys())) # noqa: SIM118 assert run['type'] == 'PRODUCTION' and run['range'] == 'DAYS_30' and run['materials_failed'] == 0 with open(f'rawData/{month}-prices.json', 'r') as f: prices = prepare.get_prices(f) prod_data, company_data = get_prod_and_company_data(db, run_id, prices) with open(f'rawData/{month}.csv', 'r', newline='') as f: data = prepare.read_data(f) prepare.write_data(month, data, prod_data, company_data) 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]]: cur = db.execute(''' SELECT ticker, entity_id, score, rank FROM leaderboard_scores JOIN materials ON materials.material_pk = leaderboard_scores.material_rowid JOIN entities ON entities.id = leaderboard_scores.entity_rowid WHERE run_id = ?''', (run_id,)) individual: dict[str, dict[str, prepare.CompanyTickerData]] = collections.defaultdict(dict) company_totals: dict[str, prepare.CompanyTotals] = collections.defaultdict(lambda: {'volume': 0.0}) universe_total_score: dict[str, float] = collections.defaultdict(int) while row := cur.fetchone(): ticker = row['ticker'] price = prices[ticker] universe_total_score[ticker] += row['score'] amount = row['score'] / 30 volume = amount * price individual[row['entity_id']][ticker] = { 'amount': amount, 'volume': volume, 'rank': row['rank'], } company_totals[row['entity_id']]['volume'] += volume prod: dict[str, prepare.ProdData] = {} # TODO: universe / 30 for ticker, total_score in universe_total_score.items(): amount = total_score / 30 prod[ticker] = {'amount': amount, 'volume': amount * prices[ticker]} company_data = {'totals': prepare.add_company_ranks(company_totals), 'individual': dict(individual)} return prod, company_data if __name__ == '__main__': main()