Source code for herostools.flow.dummy

import reprlib
from typing import Any

import numpy as np

from herostools.flow import Node, only
from herostools.helper import log




class Logger(Node):
    """Log incoming flow data and forward it unchanged.

    Every message is logged at INFO on purpose, since showing the data is the point of this node.
    """



    def _process(self, inputs: dict[str, Any], metadata: dict) -> tuple[Any, dict]:
        """Log and return incoming data.

        Args:
            inputs: Input data keyed by input name.
            metadata: Input metadata.

        Returns:
            The single input value and empty metadata.
        """
        data = only(inputs)
        log.info(f"New data from {sorted(inputs)}: {reprlib.repr(data)} with {metadata}")
        return data, {}






class AddConstant(Node):
    """Add a constant value to incoming flow data."""

    constant: float

    def __init__(self, constant: float = 5, **kwargs):
        """Initialize the node.

        Args:
            constant: Value added to incoming data.
            **kwargs: Passed on to `Node`.
        """
        super().__init__(**kwargs)
        self.constant = constant



    def _process(self, inputs: dict[str, Any], metadata: dict) -> tuple[Any, dict]:
        """Add the configured constant to incoming data.

        Args:
            inputs: Input data keyed by input name.
            metadata: Input metadata.

        Returns:
            Transformed data and empty metadata.
        """
        return only(inputs) + self.constant, {}






class Sqrt(Node):
    """Apply the square root to incoming flow data."""



    def _process(self, inputs: dict[str, Any], metadata: dict) -> tuple[Any, dict]:
        """Calculate the square root of incoming data.

        Args:
            inputs: Input data keyed by input name.
            metadata: Input metadata.

        Returns:
            Transformed data and empty metadata.
        """
        return np.sqrt(only(inputs)), {}






class ToPower(Node):
    """Raise incoming flow data to a configured exponent."""

    exponent: float

    def __init__(self, exponent: float = 2, **kwargs):
        """Initialize the node.

        Args:
            exponent: Power applied to incoming data.
            **kwargs: Passed on to `Node`.
        """
        super().__init__(**kwargs)
        self.exponent = exponent



    def _process(self, inputs: dict[str, Any], metadata: dict) -> tuple[Any, dict]:
        """Raise incoming data to the configured exponent.

        Args:
            inputs: Input data keyed by input name.
            metadata: Input metadata.

        Returns:
            Transformed data and empty metadata.
        """
        return only(inputs) ** self.exponent, {}