Source code for forml.flow._graph.port

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

Graph node port functionality.
import collections
import typing

from .. import _exception
from . import atomic  # pylint: disable=unused-import

if typing.TYPE_CHECKING:
    from forml import flow

class Type(int):
    """Input port base class."""

    def __repr__(self):
        return f'{self.__class__.__name__}[{int(self)}]'

    def __hash__(self):
        return hash(self.__class__) ^ int.__hash__(self)

    def __eq__(self, other):
        return other.__class__ is self.__class__ and int.__eq__(self, other)

class Meta(type):
    """Metaclass for singleton types."""

    def __new__(mcs, name: str, bases: tuple[type], namespace: dict[str, typing.Any]):
        value = namespace.pop('VALUE')
        instance = None

        def new(cls):
            """Injected class new method ensuring singletons with static value are only created."""
            nonlocal instance
            if instance is None:
                instance = bases[0].__new__(cls, value)
            return instance

        namespace['__new__'] = new
        return super().__new__(mcs, name, bases, namespace)

class Train(Type, metaclass=Meta):
    """Train input port."""

    VALUE = 0

class Label(Type, metaclass=Meta):
    """Label input port."""

    VALUE = 1

class Apply(Type):
    """Apply input/output port at the given index."""

class Subscription(collections.namedtuple('Subscription', 'node, port')):
    """Descriptor representing subscription node input port of given type."""

    # registry of ports subscribed on given node
    _PORTS: dict['flow.Node', set[Type]] = collections.defaultdict(set)  # TO-DO: switch to weakref

    def __new__(cls, subscriber: 'flow.Node', port: Type):
        if port in cls._PORTS[subscriber]:
            raise _exception.TopologyError('Double subscription')
        if cls._PORTS[subscriber] and (
            isinstance(port, Apply) ^ any(isinstance(s, Apply) for s in cls._PORTS[subscriber])
            raise _exception.TopologyError('Apply/Train collision')
        if isinstance(port, (Train, Label)) and any(subscriber.output):
            raise _exception.TopologyError('Publishing node trained')
        if isinstance(subscriber, atomic.Future):
            raise _exception.TopologyError('Future node subscribing')
        return super().__new__(cls, subscriber, port)

    def __repr__(self):
        return f'{self.node}@{self.port}'

    def __hash__(self):
        return hash(self.node) ^ hash(self.port)

    def __eq__(self, other: typing.Any):
        return isinstance(other, self.__class__) and self.node == other.node and self.port == other.port

    def ports(cls, subscriber: 'flow.Node') -> typing.Iterable[Type]:
        """Get subscribed ports of given atomic.

            subscriber: Node whose subscribed ports should be retrieved.

            Subscribed ports of given atomic.
        return frozenset(cls._PORTS[subscriber])

    def __del__(self):
        self._PORTS.get(self.node, {}).discard(self.port)

class Applicable:
    """Base for publisher/subscriber proxies."""

    def __init__(self, node: 'flow.Node', index: int):
        self._node: 'flow.Node' = node
        self._index: int = index

[docs]class Publishable(Applicable): """Output *Apply* port reference that can be used just for publishing.""" @property def szout(self) -> int: """Size of publisher node output. Returns: Output size. """ return self._node.szout def publish(self, subscriber: 'flow.Node', port: Type) -> None: """Publish new subscription. Args: subscriber: Node to publish to. port: Port to publish to. """ if isinstance(subscriber, atomic.Future) and subscriber is not self._node: subscriber[port].subscribe(self) return subscription = Subscription(subscriber, port) try: self.republish(subscription) except Exception as err: # TO-DO: use weakref Subscription._PORTS[subscriber].discard(port) # pylint: disable=protected-access raise err def republish(self, subscription: 'flow.Subscription') -> None: """Publish existing subscription. Args: subscription: Existing subscription descriptor. """ self._node._publish(self._index, subscription) # pylint: disable=protected-access
class Subscriptable(Applicable): """Input *Apply* port reference that can be used just for subscribing.""" @property def szin(self) -> int: """Size of publisher node input. Returns: Input size. """ return self._node.szin def subscribe(self, publisher: 'flow.Publishable') -> None: """Subscribe to the given publisher. Args: publisher: *Publishable* to subscribe to. """ publisher.publish(self._node, Apply(self._index))
[docs]class PubSub(Publishable, Subscriptable): """Input or output *Apply* port reference that can be used for both subscribing and publishing.""" @property def publisher(self) -> 'flow.Publishable': """Return just a publishable representation. Returns: Publishable *Apply* port reference. """ return Publishable(self._node, self._index) @property def subscriber(self) -> 'flow.Subscriptable': """Return just a subscriptable representation. Returns: Subscriptable *Apply* port reference. """ return Subscriptable(self._node, self._index)