piker/piker/fsp/_api.py

# piker: trading gear for hackers
# Copyright (C) Tyler Goodlet (in stewardship of 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/>.

'''
FSP (financial signal processing) apis.

'''

# TODO: things to figure the heck out:
# - how to handle non-plottable values (pyqtgraph has facility for this
#   now in `arrayToQPath()`)
# - composition of fsps / implicit chaining syntax (we need an issue)

from __future__ import annotations
from typing import (
    AsyncGenerator,
    Callable,
    NamedTuple,
    overload,
    Protocol,
)

import numpy as np
import tractor
from tractor.msg import NamespacePath

from tractor.ipc._shm import (
    ShmArray,
    NDToken,
    attach_shm_ndarray,
)
from ..data._sharedmem import maybe_open_shm_array
from ..log import get_logger

log = get_logger(__name__)

type FspHistory = dict[str, np.ndarray|None]|np.ndarray
type RealtimeValue = np.ndarray|np.number|int|float


class FspSample(NamedTuple):
    '''
    Publish all FSP fields at an absolute source coordinate.

    Unlike a field-only tick update, this record remains valid if the
    sampler has not appended its destination row yet.
    ``FspSample.fields`` are written without task checkpoints.
    This is not an atomic memory transaction for readers in other
    processes.

    '''
    index: int
    time: float
    fields: dict[str, RealtimeValue]


type FspYield = FspHistory|tuple[str, RealtimeValue]|FspSample
type FspStream = AsyncGenerator[FspYield, None]
type FspConfigValue = str|bool|int|float


# ``**P`` declares a PEP 695 ``ParamSpec`` so wrappers retain each
# decorated FSP's positional and keyword parameter types:
# https://docs.python.org/3/reference/compound_stmts.html#type-params
class FspFunc[**P](Protocol):
    @property
    def __name__(self) -> str: ...

    def __call__(
        self,
        *args: P.args,
        **kwargs: P.kwargs,

    ) -> FspStream: ...

# global fsp registry filled out by @fsp decorator below
_fsp_registry: dict[NamespacePath, Fsp] = {}


def _load_builtins() -> dict[NamespacePath, Fsp]:

    # import to implicity trigger registration via ``@fsp``
    from . import _momo  # noqa
    from . import _volume  # noqa
    from . import _trend_fsp  # noqa

    return _fsp_registry


class Fsp[**P]:
    '''
    "Financial signal processor" decorator wrapped async function.

    '''

    # TODO: checkout the advanced features from ``wrapt``:
    # - dynamic enable toggling,
    # https://wrapt.readthedocs.io/en/latest/decorators.html#dynamically-disabling-decorators
    # - custom object proxies, might be useful for implementing n-compose
    # https://wrapt.readthedocs.io/en/latest/wrappers.html#custom-object-proxies
    # - custom function wrappers,
    # https://wrapt.readthedocs.io/en/latest/wrappers.html#custom-function-wrappers

    # actor-local map of source flow shm tokens
    # + the consuming fsp *to* the consumers output
    # shm flow.
    _flow_registry: dict[
        tuple[NDToken, str],
        tuple[NDToken, ShmArray|None],
    ] = {}

    def __init__(
        self,
        func: FspFunc[P],
        *,
        outputs: tuple[str, ...] = (),
        display_name: str|None = None,
        sample_on_step: bool = False,
        event_outputs: tuple[str, ...] = (),
        **config: FspConfigValue,

    ) -> None:

        # TODO (maybe):
        # - type introspection?
        # - should we make this a wrapt object proxy?
        self.func: FspFunc[P] = func
        self.__name__ = func.__name__  # XXX: must have func-object name

        self.ns_path: NamespacePath = NamespacePath.from_ref(func)
        self.outputs = outputs
        # Snapshot processors read SHM on both quote and timer steps.
        # Sparse event fields must never be carried into a new row.
        self.sample_on_step: bool = sample_on_step
        if not set(event_outputs).issubset(outputs):
            raise ValueError('event outputs must be declared')
        self.event_outputs: tuple[str, ...] = event_outputs
        self.config: dict[str, FspConfigValue] = config

        # register with declared set.
        _fsp_registry[self.ns_path] = self

    @property
    def name(self) -> str:
        return self.__name__

    def __call__(
        self,

        # TODO: when we settle on py3.10 we should probably use the new
        # type annots from pep 612:
        # https://www.python.org/dev/peps/pep-0612/
        # instance,
        *args: P.args,
        **kwargs: P.kwargs,

    ) -> FspStream:
        return self.func(*args, **kwargs)

    def get_shm(
        self,
        src_shm: ShmArray,

    ) -> ShmArray:
        '''
        Provide access to allocated shared mem array
        for this "instance" of a signal processor for
        the given ``key``.

        The destination shm "token" and array are cached if possible to
        minimize multiple stdlib/system calls.

        '''
        dst_token, maybe_array = self._flow_registry[
            (src_shm._token, self.name)
        ]
        if maybe_array is None:
            self._flow_registry[
                (src_shm._token, self.name)
            ] = (
                dst_token,
                # "cache" the ``ShmArray`` such that
                # we call the underlying "attach" code as few
                # times as possible as per:
                # - https://github.com/pikers/piker/issues/359
                # - https://github.com/pikers/piker/issues/332
                maybe_array := attach_shm_ndarray(dst_token)
            )

        return maybe_array


@overload
def fsp[**P](
    wrapped: FspFunc[P],
    *,
    outputs: tuple[str, ...] = (),
    display_name: str|None = None,
    sample_on_step: bool = False,
    event_outputs: tuple[str, ...] = (),
    **config: FspConfigValue,

) -> Fsp[P]:
    ...


@overload
def fsp[**P](
    wrapped: None = None,
    *,
    outputs: tuple[str, ...] = (),
    display_name: str|None = None,
    sample_on_step: bool = False,
    event_outputs: tuple[str, ...] = (),
    **config: FspConfigValue,

) -> Callable[[FspFunc[P]], Fsp[P]]:
    ...


def fsp[**P](
    wrapped: FspFunc[P]|None = None,
    *,
    outputs: tuple[str, ...] = (),
    display_name: str|None = None,
    sample_on_step: bool = False,
    event_outputs: tuple[str, ...] = (),
    **config: FspConfigValue,

) -> Fsp[P]|Callable[[FspFunc[P]], Fsp[P]]:
    '''
    Wrap an FSP function directly or return its configured decorator.

    Bare ``@fsp`` calls this function with ``wrapped`` and returns an
    ``Fsp``. Parameterized ``@fsp(...)`` returns ``decorate`` first;
    Python then passes the decorated function to that closure.

    ``sample_on_step`` opts into coalesced SHM snapshot wakeups from
    quotes and sampler steps; the quote payload is not preserved.
    ``event_outputs`` names sparse fields cleared to NaN on append.

    '''

    if wrapped is None:
        def decorate(func: FspFunc[P]) -> Fsp[P]:
            return Fsp(
                func,
                outputs=outputs,
                display_name=display_name,
                sample_on_step=sample_on_step,
                event_outputs=event_outputs,
                **config,
            )

        return decorate

    return Fsp(
        wrapped,
        outputs=outputs or (wrapped.__name__,),
        display_name=display_name,
        sample_on_step=sample_on_step,
        event_outputs=event_outputs,
        **config,
    )


def maybe_mk_fsp_shm(
    sym: str,
    target: Fsp,
    size: int,
    readonly: bool = True,

) -> tuple[str, ShmArray, bool]:
    '''
    Allocate a single row shm array for an symbol-fsp pair if none
    exists, otherwise load the shm already existing for that token.

    '''
    if not isinstance(sym, str):
        raise ValueError('`sym: str` should be file-name-friendly')

    # TODO: load output types from `Fsp`
    # - should `index` be a required internal field?
    fsp_dtype = np.dtype(
        [('index', int)]
        +
        [('time', float)]
        +
        [(field_name, float) for field_name in target.outputs]
    )

    # (attempt to) uniquely key the fsp shm buffers
    # Use hash for macOS compatibility (31 char limit)
    import hashlib
    actor_name, uuid = tractor.current_actor().uid
    # Create short hash of sym and target name
    content = f'{sym}.{target.name}'
    content_hash = hashlib.md5(content.encode()).hexdigest()[:8]
    key: str = f'{uuid[:8]}_{content_hash}.fsp'

    shm, opened = maybe_open_shm_array(
        key,
        size=size,
        dtype=fsp_dtype,
        readonly=True,
    )
    return key, shm, opened