prosperity3bt 0.9.0__py3-none-any.whl → 0.10.0__py3-none-any.whl

prosperity3bt/data.py

@@ -49,6 +49,18 @@ def get_column_values(columns: list[str], indices: list[int]) -> list[int]:
     return values
 
 
+@dataclass
+class ObservationRow:
+    timestamp: int
+    bidPrice: float
+    askPrice: float
+    transportFees: float
+    exportTariff: float
+    importTariff: float
+    sugarPrice: float
+    sunlightIndex: float
+
+
 @dataclass
 class BacktestData:
     round_num: int
@@ -56,11 +68,14 @@ class BacktestData:
 
     prices: dict[int, dict[Symbol, PriceRow]]
     trades: dict[int, dict[Symbol, list[Trade]]]
+    observations: dict[int, ObservationRow]
     products: list[Symbol]
     profit_loss: dict[Symbol, float]
 
 
-def create_backtest_data(round_num: int, day_num: int, prices: list[PriceRow], trades: list[Trade]) -> BacktestData:
+def create_backtest_data(
+    round_num: int, day_num: int, prices: list[PriceRow], trades: list[Trade], observations: list[ObservationRow]
+) -> BacktestData:
     prices_by_timestamp: dict[int, dict[Symbol, PriceRow]] = defaultdict(dict)
     for row in prices:
         prices_by_timestamp[row.timestamp][row.product] = row
@@ -72,11 +87,14 @@ def create_backtest_data(round_num: int, day_num: int, prices: list[PriceRow], t
     products = sorted(set(row.product for row in prices))
     profit_loss = {product: 0.0 for product in products}
 
+    observations_by_timestamp = {row.timestamp: row for row in observations}
+
     return BacktestData(
         round_num=round_num,
         day_num=day_num,
         prices=prices_by_timestamp,
         trades=trades_by_timestamp,
+        observations=observations_by_timestamp,
         products=products,
         profit_loss=profit_loss,
     )
@@ -87,6 +105,31 @@ def has_day_data(file_reader: FileReader, round_num: int, day_num: int) -> bool:
         return file is not None
 
 
+def read_observations(file_reader: FileReader, round_num: int, day_num: int) -> list[ObservationRow]:
+    observations = []
+    with file_reader.file([f"round{round_num}", f"observations_round_{round_num}_day_{day_num}.csv"]) as file:
+        if file is None:
+            return []
+
+        for line in file.read_text(encoding="utf-8").splitlines()[1:]:
+            columns = line.split(",")
+
+            observations.append(
+                ObservationRow(
+                    timestamp=int(columns[0]),
+                    bidPrice=float(columns[1]),
+                    askPrice=float(columns[2]),
+                    transportFees=float(columns[3]),
+                    exportTariff=float(columns[4]),
+                    importTariff=float(columns[5]),
+                    sugarPrice=float(columns[6]),
+                    sunlightIndex=float(columns[7]),
+                )
+            )
+
+    return observations
+
+
 def read_day_data(file_reader: FileReader, round_num: int, day_num: int, no_names: bool) -> BacktestData:
     prices = []
     with file_reader.file([f"round{round_num}", f"prices_round_{round_num}_day_{day_num}.csv"]) as file:
@@ -134,4 +177,6 @@ def read_day_data(file_reader: FileReader, round_num: int, day_num: int, no_name
 
             break
 
-    return create_backtest_data(round_num, day_num, prices, trades)
+    observations = read_observations(file_reader, round_num, day_num)
+
+    return create_backtest_data(round_num, day_num, prices, trades, observations)