prepare.py 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  1. from __future__ import annotations
  2. import collections
  3. import dataclasses
  4. import csv
  5. import json
  6. import sys
  7. import typing
  8. def main() -> None:
  9. (month,) = sys.argv[1:]
  10. with open(f'rawData/{month}.csv', 'r', newline='') as f:
  11. data = read_data(f)
  12. bases_data: dict[str, dict[str, int]] = {r.company_id: {'bases': r.num, 'rank': r.rank} for r in data['BASES']}
  13. with open(f'www/data/base-data-{month}.json', 'w') as f:
  14. json.dump(bases_data, f)
  15. ships_data: dict[str, dict[str, int]] = {r.company_id: {'ships': r.num, 'rank': r.rank} for r in data['SHIPS']}
  16. with open(f'www/data/ship-data-{month}.json', 'w') as f:
  17. json.dump(ships_data, f)
  18. with open(f'rawData/{month}-prices.json', 'r') as f:
  19. prices = get_prices(f)
  20. prod_data, company_data = get_prod_and_company_data(data, prices)
  21. with open(f'www/data/prod-data-{month}.json', 'w') as f:
  22. json.dump(prod_data, f)
  23. with open(f'www/data/company-data-{month}.json', 'w') as f:
  24. json.dump(company_data, f)
  25. def read_data(f: typing.TextIO) -> dict[str, list[Row]]:
  26. data: dict[str, list[Row]] = collections.defaultdict(list)
  27. reader = csv.reader(f)
  28. for row in reader:
  29. data[row[0]].append(Row(int(row[1]), int(row[2]), row[3]))
  30. return data
  31. def get_prices(f: typing.TextIO) -> typing.Mapping[str, float]:
  32. raw_prices: typing.Sequence[Price] = json.load(f)
  33. volumes: dict[str, float] = collections.defaultdict(float)
  34. traded: dict[str, int] = collections.defaultdict(int)
  35. for price in raw_prices:
  36. if price['Traded30D'] is None:
  37. continue
  38. assert price['VWAP30D'] is not None
  39. volumes[price['MaterialTicker']] += price['VWAP30D'] * price['Traded30D']
  40. traded[price['MaterialTicker']] += price['Traded30D']
  41. prices = {ticker: volume / traded[ticker] for ticker, volume in volumes.items()}
  42. hardcoded_prices = {
  43. 'AFP': 116868,
  44. 'ANZ': 70601,
  45. 'BFP': 23408,
  46. 'BND': 230,
  47. 'BID': 47011,
  48. 'CRU': 169623,
  49. 'CQT': 378452,
  50. 'FUN': 124010,
  51. 'GCH': 18303,
  52. 'GNZ': 30361,
  53. 'HNZ': 93580,
  54. 'PFG': 2677222,
  55. 'PK': 869,
  56. 'RDS': 598170,
  57. 'SDM': 1721027,
  58. 'SST': 5863587,
  59. 'SU': 157860,
  60. 'TOR': 540169,
  61. 'VCB': 673713,
  62. 'WOR': 202000,
  63. }
  64. assert frozenset(prices).isdisjoint(hardcoded_prices)
  65. prices.update(hardcoded_prices)
  66. return prices
  67. def get_prod_and_company_data(data: dict[str, list[Row]], prices: typing.Mapping[str, float]
  68. ) -> tuple[typing.Mapping[str, ProdData], typing.Mapping[str, typing.Any]]:
  69. prod: dict[str, ProdData] = {}
  70. individual: dict[str, dict[str, CompanyTickerData]] = collections.defaultdict(dict)
  71. totals: dict[str, CompanyTotals] = collections.defaultdict(lambda: {'volume': 0.0})
  72. for section, rows in data.items():
  73. if (ticker := get_production_ticker(section)) is None:
  74. continue
  75. price = prices.get(ticker)
  76. if price is None:
  77. continue
  78. prod_amount = sum(row.num for row in rows) / 30
  79. prod[ticker] = {'amount': prod_amount, 'volume': prod_amount * price}
  80. for row in rows:
  81. amount = row.num / 30
  82. volume = amount * price
  83. individual[row.company_id][ticker] = {
  84. 'amount': amount,
  85. 'volume': volume,
  86. 'rank': row.rank,
  87. }
  88. totals[row.company_id]['volume'] += volume
  89. company_data = {'totals': add_company_ranks(totals), 'individual': dict(individual)}
  90. return prod, company_data
  91. def get_production_ticker(section: str) -> str | None:
  92. prefix = 'PRODUCTION_'
  93. suffix = '_DAYS_30'
  94. if not section.startswith(prefix) or not section.endswith(suffix):
  95. return None
  96. return section[len(prefix):-len(suffix)]
  97. def add_company_ranks(totals: dict[str, CompanyTotals]) -> dict[str, CompanyTotals]:
  98. ranked = sorted(totals.items(), key=lambda item: item[1]['volume'], reverse=True)
  99. for rank, (company_id, company_totals) in enumerate(ranked, start=1):
  100. company_totals['volumeRank'] = rank
  101. return totals
  102. @dataclasses.dataclass(frozen=True, slots=True, eq=False)
  103. class Row:
  104. rank: int
  105. num: int
  106. company_id: str
  107. class Price(typing.TypedDict):
  108. MaterialTicker: str
  109. VWAP30D: float | None
  110. Traded30D: int | None
  111. class ProdData(typing.TypedDict):
  112. amount: float
  113. volume: float
  114. class CompanyTickerData(typing.TypedDict):
  115. amount: float
  116. volume: float
  117. rank: int
  118. class CompanyTotals(typing.TypedDict, total=False):
  119. volume: float
  120. volumeRank: int
  121. if __name__ == '__main__':
  122. main()