piker/piker/ui/_remote_ctl.py

884 lines
29 KiB
Python

# piker: trading gear for hackers
# Copyright (C) Tyler Goodlet (in stewardship for pikers)
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
'''
Remote control tasks for sending annotations (and maybe more cmds) to
a chart from some other actor.
'''
from __future__ import annotations
from contextlib import (
asynccontextmanager as acm,
AsyncExitStack,
)
from functools import partial
from pprint import pformat
from typing import (
AsyncContextManager,
Literal,
)
from uuid import uuid4
import pyqtgraph as pg
import tractor
import trio
from tractor import trionics
from tractor import (
Portal,
Context,
MsgStream,
)
from piker.log import get_logger
from piker.types import Struct
from piker.service import find_service
from piker.brokers import SymbolNotFound
from piker.ui.qt import (
QGraphicsItem,
)
from PyQt6.QtGui import QFont
from ._display import DisplayState
from ._interaction import ChartView
from ._editors import (
SelectRect,
ArrowEditor,
)
from ._chart import ChartPlotWidget
from ._dataviz import Viz
from ._style import hcolor
log = get_logger(__name__)
# NOTE: this is UPDATED by the `._display.graphics_update_loop()`
# once all chart widgets / Viz per flume have been initialized
# allowing for remote annotation (control) of any chart-actor's mkt
# feed by fqme lookup Bo
_dss: dict[str, DisplayState] = {}
# stash each and every client connection so that they can all
# be cancelled on shutdown/error.
# TODO: make `tractor.Context` hashable via is `.cid: str`?
# _ctxs: set[Context] = set()
# TODO: use type statements from 3.12+
IpcCtxTable = dict[
str, # each `Context.cid`
tuple[
Context, # handle for ctx-cancellation
set[int] # set of annotation (instance) ids
]
]
_ctxs: IpcCtxTable = {}
# XXX: global map of all uniquely created annotation-graphics so
# that they can be mutated (eventually) by a client.
# NOTE: this map is only populated on the `chart` actor side (aka
# the "annotations server" which actually renders to a Qt canvas).
# type AnnotsTable = dict[int, QGraphicsItem]
AnnotsTable = dict[int, QGraphicsItem]
EditorsTable = dict[int, ArrowEditor]
_annots: AnnotsTable = {}
_editors: EditorsTable = {}
def rm_annot(
annot: ArrowEditor|SelectRect|pg.TextItem
) -> bool:
global _editors
match annot:
case pg.ArrowItem():
editor = _editors[annot._uid]
editor.remove(annot)
# ^TODO? only remove each arrow or all?
# if editor._arrows:
# editor.remove_all()
# else:
# log.warning(
# f'Annot already removed!\n'
# f'{annot!r}\n'
# )
return True
case SelectRect():
annot.delete()
return True
case pg.TextItem():
scene = annot.scene()
if scene:
scene.removeItem(annot)
return True
return False
async def serve_rc_annots(
ipc_key: str,
annot_req_stream: MsgStream,
dss: dict[str, DisplayState],
ctxs: IpcCtxTable,
annots: AnnotsTable,
) -> None:
'''
A small viz(ualization) server for remote ctl of chart
annotations.
'''
global _editors
async for msg in annot_req_stream:
match msg:
case {
'cmd': 'SelectRect',
'fqme': fqme,
'timeframe': timeframe,
'meth': str(meth),
'kwargs': dict(kwargs),
}:
ds: DisplayState = _dss[fqme]
try:
chart: ChartPlotWidget = {
60: ds.hist_chart,
1: ds.chart,
}[timeframe]
except KeyError:
msg: str = (
f'No chart for timeframe={timeframe}s, '
f'skipping rect annotation'
)
log.exeception(msg)
await annot_req_stream.send({'error': msg})
continue
cv: ChartView = chart.cv
# NEW: if timestamps provided, lookup current indices
# from shm to ensure alignment with current buffer
# state
start_time = kwargs.pop('start_time', None)
end_time = kwargs.pop('end_time', None)
if (
start_time is not None
and end_time is not None
):
viz: Viz = chart.get_viz(fqme)
shm = viz.shm
arr = shm.array
# lookup start index
start_matches = arr[arr['time'] == start_time]
if len(start_matches) == 0:
msg: str = (
f'No shm entry for start_time={start_time}, '
f'skipping rect'
)
log.error(msg)
await annot_req_stream.send({'error': msg})
continue
# lookup end index
end_matches = arr[arr['time'] == end_time]
if len(end_matches) == 0:
msg: str = (
f'No shm entry for end_time={end_time}, '
f'skipping rect'
)
log.error(msg)
await annot_req_stream.send({'error': msg})
continue
# get close price from start bar, open from end
# bar
start_idx = float(start_matches[0]['index'])
end_idx = float(end_matches[0]['index'])
start_close = float(start_matches[0]['close'])
end_open = float(end_matches[0]['open'])
# reconstruct start_pos and end_pos with
# looked-up indices
from_idx: float = 0.16 - 0.06 # BGM offset
kwargs['start_pos'] = (
start_idx + 1 - from_idx,
start_close,
)
kwargs['end_pos'] = (
end_idx + from_idx,
end_open,
)
# annot type lookup from cmd
rect = SelectRect(
viewbox=cv,
# TODO: make this more dynamic?
# -[ ] pull from conf.toml?
# -[ ] add `.set_color()` method to type?
# -[ ] make a green/red based on direction
# instead of default static color?
color=kwargs.pop('color', None),
)
# XXX NOTE: this is REQUIRED to set the rect
# resize callback!
rect.chart: ChartPlotWidget = chart
# delegate generically to the requested method
getattr(rect, meth)(**kwargs)
rect.show()
# XXX: store absolute coords for repositioning
# during viz redraws (eg backfill updates)
rect._meth = meth
rect._kwargs = kwargs
aid: int = id(rect)
annots[aid] = rect
aids: set[int] = ctxs[ipc_key][1]
aids.add(aid)
await annot_req_stream.send(aid)
case {
'cmd': 'ArrowEditor',
'fqme': fqme,
'timeframe': timeframe,
'meth': 'add'|'remove' as meth,
'kwargs': {
'x': float(x),
'y': float(y),
'pointing': pointing,
'color': color,
'aid': str()|None as aid,
'alpha': int(alpha),
'headLen': int()|float()|None as headLen,
'headWidth': int()|float()|None as headWidth,
'tailLen': int()|float()|None as tailLen,
'tailWidth': int()|float()|None as tailWidth,
'pxMode': bool(pxMode),
'time': int()|float()|None as timestamp,
},
# ?TODO? split based on method fn-sigs?
# 'pointing',
}:
ds: DisplayState = _dss[fqme]
try:
chart: ChartPlotWidget = {
60: ds.hist_chart,
1: ds.chart,
}[timeframe]
except KeyError:
log.warning(
f'No chart for timeframe={timeframe}s, '
f'skipping arrow annotation'
)
# return -1 to indicate failure
await annot_req_stream.send(-1)
continue
cv: ChartView = chart.cv
godw = chart.linked.godwidget
# NEW: if timestamp provided, lookup current index
# from shm to ensure alignment with current buffer
# state
if timestamp is not None:
viz: Viz = chart.get_viz(fqme)
shm = viz.shm
arr = shm.array
# find index where time matches timestamp
matches = arr[arr['time'] == timestamp]
if len(matches) == 0:
log.error(
f'No shm entry for timestamp={timestamp}, '
f'skipping arrow annotation'
)
await annot_req_stream.send(-1)
continue
# use the matched row's index as x
x = float(matches[0]['index'])
arrows = ArrowEditor(godw=godw)
# `.add/.remove()` API
if meth != 'add':
# await tractor.pause()
raise ValueError(
f'Invalid arrow-edit request ?\n'
f'{msg!r}\n'
)
aid: str = str(uuid4())
arrow: pg.ArrowItem = arrows.add(
plot=chart.plotItem,
uid=aid,
x=x,
y=y,
pointing=pointing,
color=color,
alpha=alpha,
headLen=headLen,
headWidth=headWidth,
tailLen=tailLen,
tailWidth=tailWidth,
pxMode=pxMode,
)
# XXX: store absolute coords for repositioning
# during viz redraws (eg backfill updates)
arrow._abs_x = x
arrow._abs_y = y
annots[aid] = arrow
_editors[aid] = arrows
aids: set[int] = ctxs[ipc_key][1]
aids.add(aid)
await annot_req_stream.send(aid)
case {
'cmd': 'TextItem',
'fqme': fqme,
'timeframe': timeframe,
'kwargs': {
'text': str(text),
'x': int()|float() as x,
'y': int()|float() as y,
'color': color,
'anchor': list(anchor),
'font_size': int()|None as font_size,
'time': int()|float()|None as timestamp,
},
}:
ds: DisplayState = _dss[fqme]
try:
chart: ChartPlotWidget = {
60: ds.hist_chart,
1: ds.chart,
}[timeframe]
except KeyError:
log.warning(
f'No chart for timeframe={timeframe}s, '
f'skipping text annotation'
)
await annot_req_stream.send(-1)
continue
# NEW: if timestamp provided, lookup current index
# from shm to ensure alignment with current buffer
# state
if timestamp is not None:
viz: Viz = chart.get_viz(fqme)
shm = viz.shm
arr = shm.array
# find index where time matches timestamp
matches = arr[arr['time'] == timestamp]
if len(matches) == 0:
log.error(
f'No shm entry for timestamp={timestamp}, '
f'skipping text annotation'
)
await annot_req_stream.send(-1)
continue
# use the matched row's index as x, +1 for text
# offset
x = float(matches[0]['index']) + 1
# convert named color to hex
color_hex: str = hcolor(color)
# create text item
text_item: pg.TextItem = pg.TextItem(
text=text,
color=color_hex,
anchor=anchor,
# ?TODO, pin to github:main for this?
# legacy, can have scaling ish?
# ensureInBounds=True,
)
# apply font size (default to DpiAwareFont if not
# provided)
if font_size is None:
from ._style import get_fonts
font, font_small = get_fonts()
font_size = font_small.px_size - 1
qfont: QFont = text_item.textItem.font()
qfont.setPixelSize(font_size)
text_item.setFont(qfont)
text_item.setPos(x, y)
chart.plotItem.addItem(text_item)
# XXX: store absolute coords for repositioning
# during viz redraws (eg backfill updates)
text_item._abs_x = x
text_item._abs_y = y
aid: str = str(uuid4())
annots[aid] = text_item
aids: set[int] = ctxs[ipc_key][1]
aids.add(aid)
await annot_req_stream.send(aid)
case {
'cmd': 'remove',
'aid': int(aid)|str(aid),
}:
# NOTE: this is normally entered on
# a client's annotation de-alloc normally
# prior to detach or modify.
annot: QGraphicsItem = annots[aid]
assert rm_annot(annot)
# respond to client indicating annot
# was indeed deleted.
await annot_req_stream.send(aid)
case {
'cmd': 'redraw',
'fqme': fqme,
'timeframe': timeframe,
# TODO: maybe more fields?
# 'render': int(aid),
# 'viz_name': str(viz_name),
}:
# NOTE: old match from the 60s display loop task
# | {
# 'backfilling': (str(viz_name), timeframe),
# }:
ds: DisplayState = _dss[fqme]
viz: Viz = {
60: ds.hist_viz,
1: ds.viz,
}[timeframe]
log.warning(
f'Forcing VIZ REDRAW:\n'
f'fqme: {fqme}\n'
f'timeframe: {timeframe}\n'
)
viz.reset_graphics()
# XXX: reposition all annotations to ensure they
# stay aligned with viz data after reset (eg during
# backfill when abs-index range changes)
n_repositioned: int = 0
for aid, annot in annots.items():
# arrows and text items use abs x,y coords
if (
hasattr(annot, '_abs_x')
and
hasattr(annot, '_abs_y')
):
annot.setPos(
annot._abs_x,
annot._abs_y,
)
n_repositioned += 1
# rects use method + kwargs
elif (
hasattr(annot, '_meth')
and
hasattr(annot, '_kwargs')
):
getattr(annot, annot._meth)(**annot._kwargs)
n_repositioned += 1
if n_repositioned:
log.info(
f'Repositioned {n_repositioned} annotation(s) '
f'after viz redraw'
)
case _:
log.error(
'Unknown remote annotation cmd:\n'
f'{pformat(msg)}'
)
@tractor.context
async def remote_annotate(
ctx: Context,
) -> None:
global _dss, _ctxs
if not _dss:
raise RuntimeError(
'Race condition on chart-init state ??\n'
'Anoter actor is trying to annoate this chart '
'before it has fully spawned.\n'
)
assert _dss
_ctxs[ctx.cid] = (ctx, set())
# send back full fqme symbology to caller
await ctx.started(list(_dss))
# open annot request handler stream
async with ctx.open_stream() as annot_req_stream:
try:
await serve_rc_annots(
ipc_key=ctx.cid,
annot_req_stream=annot_req_stream,
dss=_dss,
ctxs=_ctxs,
annots=_annots,
)
finally:
# ensure all annots for this connection are deleted
# on any final teardown
(_ctx, aids) = _ctxs[ctx.cid]
assert _ctx is ctx
for aid in aids:
annot: QGraphicsItem = _annots[aid]
assert rm_annot(annot)
class AnnotCtl(Struct):
'''
A control for remote "data annotations".
You know those "squares they always show in machine vision
UIs.." this API allows you to remotely control stuff like that
in some other graphics actor.
'''
ctx2fqmes: dict[str, str]
fqme2ipc: dict[str, MsgStream]
_annot_stack: AsyncExitStack
# runtime-populated mapping of all annotation
# ids to their equivalent IPC msg-streams.
_ipcs: dict[int, MsgStream] = {}
def _get_ipc(
self,
fqme: str,
) -> MsgStream:
ipc: MsgStream = self.fqme2ipc.get(fqme)
if ipc is None:
raise SymbolNotFound(
'No chart (actor) seems to have mkt feed loaded?\n'
f'{fqme}'
)
return ipc
async def add_rect(
self,
fqme: str,
timeframe: float,
start_pos: tuple[float, float],
end_pos: tuple[float, float],
# TODO: a `Literal['view', 'scene']` for this?
domain: str = 'view', # or 'scene'
color: str = 'dad_blue',
from_acm: bool = False,
# NEW: optional timestamps for server-side index lookup
start_time: float|None = None,
end_time: float|None = None,
) -> int|None:
'''
Add a `SelectRect` annotation to the target view, return
the instances `id(obj)` from the remote UI actor.
'''
ipc: MsgStream = self._get_ipc(fqme)
with trio.fail_after(3):
await ipc.send({
'fqme': fqme,
'cmd': 'SelectRect',
'timeframe': timeframe,
# 'meth': str(meth),
'meth': 'set_view_pos' if domain == 'view' else 'set_scene_pos',
'kwargs': {
'start_pos': tuple(start_pos),
'end_pos': tuple(end_pos),
'color': color,
'update_label': False,
'start_time': start_time,
'end_time': end_time,
},
})
aid: int|dict = await ipc.receive()
match aid:
case {'error': str(msg)}:
log.error(msg)
return None
self._ipcs[aid] = ipc
if not from_acm:
self._annot_stack.push_async_callback(
partial(
self.remove,
aid,
)
)
return aid
async def remove(
self,
aid: int,
) -> bool:
'''
Remove an existing annotation by instance id.
'''
ipc: MsgStream = self._ipcs[aid]
await ipc.send({
'cmd': 'remove',
'aid': aid,
})
removed: bool = await ipc.receive()
return removed
@acm
async def open_rect(
self,
**kwargs,
) -> int:
try:
aid: int = await self.add_rect(
from_acm=True,
**kwargs,
)
yield aid
finally:
# async ipc send op
with trio.CancelScope(shield=True):
await self.remove(aid)
async def redraw(
self,
fqme: str,
timeframe: float,
) -> None:
await self._get_ipc(fqme).send({
'cmd': 'redraw',
'fqme': fqme,
# 'render': int(aid),
# 'viz_name': str(viz_name),
'timeframe': timeframe,
})
async def add_arrow(
self,
fqme: str,
timeframe: float,
x: float,
y: float,
pointing: Literal[
'up',
'down',
],
# TODO: a `Literal['view', 'scene']` for this?
# domain: str = 'view', # or 'scene'
color: str = 'dad_blue',
alpha: int = 116,
headLen: float|None = None,
headWidth: float|None = None,
tailLen: float|None = None,
tailWidth: float|None = None,
pxMode: bool = True,
from_acm: bool = False,
# NEW: optional timestamp for server-side index lookup
time: float|None = None,
) -> int|None:
'''
Add a `SelectRect` annotation to the target view, return
the instances `id(obj)` from the remote UI actor.
'''
ipc: MsgStream = self._get_ipc(fqme)
with trio.fail_after(3):
await ipc.send({
'fqme': fqme,
'cmd': 'ArrowEditor',
'timeframe': timeframe,
# 'meth': str(meth),
'meth': 'add',
'kwargs': {
'x': float(x),
'y': float(y),
'color': color,
'pointing': pointing, # up|down
'alpha': alpha,
'aid': None,
'headLen': headLen,
'headWidth': headWidth,
'tailLen': tailLen,
'tailWidth': tailWidth,
'pxMode': pxMode,
'time': time, # for server-side index lookup
},
})
aid: int|dict = await ipc.receive()
match aid:
case {'error': str(msg)}:
log.error(msg)
return None
self._ipcs[aid] = ipc
if not from_acm:
self._annot_stack.push_async_callback(
partial(
self.remove,
aid,
)
)
return aid
async def add_text(
self,
fqme: str,
timeframe: float,
text: str,
x: float,
y: float,
color: str|tuple = 'dad_blue',
anchor: tuple[float, float] = (0, 1),
font_size: int|None = None,
from_acm: bool = False,
# NEW: optional timestamp for server-side index lookup
time: float|None = None,
) -> int|None:
'''
Add a `pg.TextItem` annotation to the target view.
anchor: (x, y) where (0,0) is upper-left, (1,1) is lower-right
font_size: pixel size for font, defaults to `_font.font.pixelSize()`
'''
ipc: MsgStream = self._get_ipc(fqme)
with trio.fail_after(3):
await ipc.send({
'fqme': fqme,
'cmd': 'TextItem',
'timeframe': timeframe,
'kwargs': {
'text': text,
'x': float(x),
'y': float(y),
'color': color,
'anchor': tuple(anchor),
'font_size': font_size,
'time': time, # for server-side index lookup
},
})
aid: int|dict = await ipc.receive()
match aid:
case {'error': str(msg)}:
log.error(msg)
return None
self._ipcs[aid] = ipc
if not from_acm:
self._annot_stack.push_async_callback(
partial(
self.remove,
aid,
)
)
return aid
@acm
async def open_annot_ctl(
uid: tuple[str, str] | None = None,
) -> AnnotCtl:
# TODO: load connetion to a specific chart actor
# -[ ] pull from either service scan or config
# -[ ] return some kinda client/proxy thinger?
# -[ ] maybe we should finally just provide this as
# a `tractor.hilevel.CallableProxy` or wtv?
# -[ ] use this from the storage.cli stuff to mark up gaps!
maybe_portals: list[Portal] | None
fqmes: list[str]
async with find_service(
service_name='chart',
first_only=False,
) as maybe_portals:
ctx_mngrs: list[AsyncContextManager] = []
# TODO: print the current discoverable actor UID set
# here as well?
if not maybe_portals:
raise RuntimeError(
'No chart actors found in service domain?'
)
for portal in maybe_portals:
ctx_mngrs.append(
portal.open_context(remote_annotate)
)
ctx2fqmes: dict[str, set[str]] = {}
fqme2ipc: dict[str, MsgStream] = {}
stream_ctxs: list[AsyncContextManager] = []
async with (
trionics.gather_contexts(ctx_mngrs) as ctxs,
):
for (ctx, fqmes) in ctxs:
stream_ctxs.append(ctx.open_stream())
# fill lookup table of mkt addrs to IPC ctxs
for fqme in fqmes:
if other := fqme2ipc.get(fqme):
raise ValueError(
f'More then one chart displays {fqme}!?\n'
'Other UI actor info:\n'
f'channel: {other._ctx.chan}]\n'
f'actor uid: {other._ctx.chan.uid}]\n'
f'ctx id: {other._ctx.cid}]\n'
)
ctx2fqmes.setdefault(
ctx.cid,
set(),
).add(fqme)
async with trionics.gather_contexts(stream_ctxs) as streams:
for stream in streams:
fqmes: set[str] = ctx2fqmes[stream._ctx.cid]
for fqme in fqmes:
fqme2ipc[fqme] = stream
# NOTE: on graceful teardown we always attempt to
# remove all annots that were created by the
# entering client.
# TODO: should we maybe instead/also do this on the
# server-actor side so that when a client
# disconnects we always delete all annotations by
# default instaead of expecting the client to?
async with AsyncExitStack() as annots_stack:
client = AnnotCtl(
ctx2fqmes=ctx2fqmes,
fqme2ipc=fqme2ipc,
_annot_stack=annots_stack,
)
yield client