117 lines
2.8 KiB
Python
117 lines
2.8 KiB
Python
from typing import (
|
|
Annotated,
|
|
AsyncGenerator,
|
|
Final,
|
|
Generator,
|
|
List,
|
|
Literal,
|
|
Optional,
|
|
Tuple,
|
|
TypeVar,
|
|
TypedDict,
|
|
Any,
|
|
cast,
|
|
)
|
|
|
|
from loguru import logger
|
|
import numpy as np
|
|
import plotly.graph_objects as go
|
|
import streamlit as st
|
|
import anyio
|
|
from anyio.abc import TaskGroup, UDPSocket
|
|
from anyio import create_memory_object_stream, create_udp_socket
|
|
from anyio.streams.memory import MemoryObjectSendStream, MemoryObjectReceiveStream
|
|
from threading import Thread
|
|
from time import sleep
|
|
from pydantic import BaseModel, computed_field
|
|
from datetime import datetime
|
|
import awkward as ak
|
|
from awkward import Array as AwkwardArray, Record as AwkwardRecord
|
|
from app.model import AlgoReport
|
|
from collections import deque
|
|
|
|
|
|
# https://handmadesoftware.medium.com/streamlit-asyncio-and-mongodb-f85f77aea825
|
|
class AppState(TypedDict):
|
|
worker_thread: Thread
|
|
message_queue: MemoryObjectReceiveStream[bytes]
|
|
task_group: TaskGroup
|
|
history: deque[AlgoReport]
|
|
|
|
|
|
UDP_LISTEN_PORT: Final[int] = 50_000
|
|
MAX_LENGTH = 600
|
|
NDArray = np.ndarray
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def unwrap(value: Optional[T]) -> T:
|
|
if value is None:
|
|
raise ValueError("Value is None")
|
|
return value
|
|
|
|
|
|
@st.cache_resource
|
|
def resource(params: Any = None):
|
|
set_ev = anyio.Event()
|
|
tx, rx = create_memory_object_stream[bytes]()
|
|
tg: Optional[TaskGroup] = None
|
|
|
|
async def poll_task():
|
|
nonlocal set_ev
|
|
nonlocal tg
|
|
tg = anyio.create_task_group()
|
|
set_ev.set()
|
|
async with tg:
|
|
async with await create_udp_socket(
|
|
local_port=UDP_LISTEN_PORT, reuse_port=True
|
|
) as udp:
|
|
async for packet, _ in udp:
|
|
await tx.send(packet)
|
|
|
|
tr = Thread(target=anyio.run, args=(poll_task,))
|
|
tr.start()
|
|
while not set_ev.is_set():
|
|
sleep(0.01)
|
|
logger.info("Poll task initialized")
|
|
state: AppState = {
|
|
"worker_thread": tr,
|
|
"message_queue": rx,
|
|
"task_group": unwrap(tg),
|
|
"history": deque(maxlen=MAX_LENGTH),
|
|
}
|
|
return state
|
|
|
|
|
|
def main():
|
|
state = resource()
|
|
logger.info("Resource created")
|
|
history = state["history"]
|
|
|
|
def on_export():
|
|
raise NotImplementedError
|
|
|
|
def on_clear():
|
|
raise NotImplementedError
|
|
|
|
st.button(
|
|
"Export", help="Export the current data to a parquet file", on_click=on_export
|
|
)
|
|
st.button("Clear", help="Clear the current data", on_click=on_clear)
|
|
pannel = st.empty()
|
|
while True:
|
|
try:
|
|
message = state["message_queue"].receive_nowait()
|
|
except anyio.WouldBlock:
|
|
continue
|
|
# TODO: plot
|
|
# fig = go.Figure(scatters)
|
|
# pannel.plotly_chart(fig)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
|
# 1659A202
|