31 lines
956 B
Python
31 lines
956 B
Python
from collections import defaultdict, deque
|
|
from typing import Protocol
|
|
|
|
from app.bars.models import Bar, Timeframe
|
|
|
|
|
|
class BarStore(Protocol):
|
|
def put(self, bar: Bar) -> None: ...
|
|
|
|
def get(self, tf: Timeframe, limit: int | None = None) -> list[Bar]: ...
|
|
|
|
|
|
class InMemoryBarStore:
|
|
def __init__(self, max_bars_per_tf: int = 5000):
|
|
self._bars: dict[Timeframe, deque[Bar]] = defaultdict(
|
|
lambda: deque(maxlen=max_bars_per_tf)
|
|
)
|
|
|
|
def put(self, bar: Bar) -> None:
|
|
bars = self._bars[bar.tf]
|
|
if bars and bars[-1].t == bar.t:
|
|
bars[-1] = bar
|
|
elif not bars or bar.t > bars[-1].t:
|
|
bars.append(bar)
|
|
|
|
def get(self, tf: Timeframe, limit: int | None = None) -> list[Bar]:
|
|
bars = list(self._bars[tf])
|
|
return bars[-limit:] if limit is not None else bars
|
|
|
|
def counts(self) -> dict[str, int]:
|
|
return {tf.value: len(self._bars[tf]) for tf in Timeframe}
|