stable/freqtrade/rpc/api_server/ws/schema.py
2022-09-07 15:08:01 -06:00

79 lines
1.9 KiB
Python

from datetime import datetime
from typing import Any, Dict, List, Optional
from pandas import DataFrame
from pydantic import BaseModel, ValidationError
from freqtrade.constants import PairWithTimeframe
from freqtrade.enums.rpcmessagetype import RPCMessageType, RPCRequestType
__all__ = ('WSRequestSchema', 'WSMessageSchema', 'ValidationError')
class BaseArbitraryModel(BaseModel):
class Config:
arbitrary_types_allowed = True
class WSRequestSchema(BaseArbitraryModel):
type: RPCRequestType
data: Optional[Any] = None
class WSMessageSchema(BaseArbitraryModel):
type: RPCMessageType
data: Optional[Any] = None
class Config:
extra = 'allow'
# ------------------------------ REQUEST SCHEMAS ----------------------------
class WSSubscribeRequest(WSRequestSchema):
type: RPCRequestType = RPCRequestType.SUBSCRIBE
data: List[RPCMessageType]
class WSWhitelistRequest(WSRequestSchema):
type: RPCRequestType = RPCRequestType.WHITELIST
data: None = None
class WSAnalyzedDFRequest(WSRequestSchema):
type: RPCRequestType = RPCRequestType.ANALYZED_DF
data: Dict[str, Any] = {"limit": 1500}
# ------------------------------ MESSAGE SCHEMAS ----------------------------
class WSWhitelistMessage(WSMessageSchema):
type: RPCMessageType = RPCMessageType.WHITELIST
data: List[str]
class WSAnalyzedDFMessage(WSMessageSchema):
class AnalyzedDFData(BaseArbitraryModel):
key: PairWithTimeframe
df: DataFrame
la: datetime
type: RPCMessageType = RPCMessageType.ANALYZED_DF
data: AnalyzedDFData
# --------------------------------------------------------------------------
if __name__ == "__main__":
message = WSAnalyzedDFMessage(
data={
"key": ("1", "5m", "spot"),
"df": DataFrame(),
"la": datetime.now()
}
)
print(message)