Prechádzať zdrojové kódy

prepare_from_sqlite.py

raylu 3 týždňov pred
rodič
commit
0a0d2d0161
1 zmenil súbory, kde vykonal 69 pridanie a 0 odobranie
  1. 69 0
      py/prepare_from_sqlite.py

+ 69 - 0
py/prepare_from_sqlite.py

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