Toggle Light / Dark / Auto color theme

Toggle table of contents sidebar

Source code for jina.enums

"""
Miscellaneous enums used in Jina.

To use these enums in YAML config, following the example below:

.. highlight:: yaml
.. code-block:: yaml

      chunk_idx:
        uses: index/chunk.yml
        parallel: ${{PARALLEL}}
        parallel_type: !PollingType ANY
        # or
        parallel_type: ANY
        # or
        parallel_type: any
"""

from enum import Enum, IntEnum, EnumMeta


class EnumType(EnumMeta):
    """The metaclass for BetterEnum."""

    def __new__(cls, *args, **kwargs):
        """Register a new EnumType

        :param args: args passed to super()
        :param kwargs: kwargs passed to super()
        :return: the registry class
        """
        _cls = super().__new__(cls, *args, **kwargs)
        return cls.register_class(_cls)

    @staticmethod
    def register_class(cls):
        """
        Register the class for dumping loading.

        :param cls: Target class.
        :return: Registered class.
        """
        reg_cls_set = getattr(cls, '_registered_class', set())
        if cls.__name__ not in reg_cls_set:
            reg_cls_set.add(cls.__name__)
            setattr(cls, '_registered_class', reg_cls_set)
        from jina.jaml import JAML

        JAML.register(cls)
        return cls


class BetterEnum(IntEnum, metaclass=EnumType):
    """The base class of Enum used in Jina."""

    def __str__(self):
        return self.name

    @classmethod
    def from_string(cls, s: str):
        """
        Parse the enum from a string.

        :param s: string representation of the enum value
        :return: enum value
        """
        try:
            return cls[s.upper()]
        except KeyError:
            raise ValueError(
                f'{s.upper()} is not a valid enum for {cls!r}, must be one of {list(cls)}'
            )

    @classmethod
    def _to_yaml(cls, representer, data):
        """Required by :mod:`pyyaml`.

        .. note::
            In principle, this should inherit from :class:`JAMLCompatible` directly,
            however, this method is too simple and thus replaced the parent method.
        :param representer: pyyaml representer
        :param data: enum value
        :return: yaml representation
        """
        return representer.represent_scalar(
            'tag:yaml.org,2002:str', str(data), style='"'
        )

    @classmethod
    def _from_yaml(cls, constructor, node):
        """Required by :mod:`pyyaml`.

        .. note::
            In principle, this should inherit from :class:`JAMLCompatible` directly,
            however, this method is too simple and thus replaced the parent method.
        :param constructor: unused
        :param node: node to derive the enum value from
        :return: enum value
        """
        return cls.from_string(node.value)


class PollingType(BetterEnum):
    """The enum for representing the parallel type of pods in a deployment."""

    ANY = 1  #: one of the shards will receive the message
    ALL = 2  #: all shards will receive the message, blocked until all done with the message
    ALL_ASYNC = 3  #: (reserved) all replica will receive the message, but any one of them can return, useful in backup

    @property
    def is_push(self) -> bool:
        """
        Check if :class:`PollingType` is using `push` protocol.

        :return: True if this :class:`PollingType` is using `push` protocol else False.
        """
        return self.value == 1

    @property
    def is_block(self) -> bool:
        """
        Check if :class:`PollingType` is using `block` protocol.

        :return: True if this :class:`PollingType` is requiring `block` protocol else False.
        """
        return self.value == 2


class LogVerbosity(BetterEnum):
    """Verbosity level of the logger."""

    DEBUG = 10
    INFO = 20
    SUCCESS = 25
    WARNING = 30
    ERROR = 40
    CRITICAL = 50


class FlowBuildLevel(BetterEnum):
    """
    The enum for representing a flow's build level.

    Some :class:`jina.orchestrate.flow.Flow` class functions require certain build level to run.
    """

    EMPTY = 0  #: Nothing is built
    GRAPH = 1  #: The underlying graph is built, you may visualize the flow
    RUNNING = 2  #: the graph is started and all deployment are running


class GatewayProtocolType(BetterEnum):
    """
    Gateway communication protocol
    """

    GRPC = 0
    HTTP = 1
    WEBSOCKET = 2


class PodRoleType(BetterEnum):
    """The enum of a Pod role."""

    HEAD = 0
    WORKER = 1
    GATEWAY = 2


class DeploymentRoleType(BetterEnum):
    """The enum of a Deploymen role for visualization."""

    DEPLOYMENT = 0
    JOIN = 1
    INSPECT = 2
    GATEWAY = 3
    INSPECT_AUX_PASS = 4
    JOIN_INSPECT = 5

    @property
    def is_inspect(self) -> bool:
        """
        If the role is inspect deployment related.

        :return: True if the Deployment role is inspect related else False.
        """
        return self.value in {2, 4}


class RequestType(BetterEnum):
    """The enum of Client mode."""

    DATA = 0
    CONTROL = 1


class CompressAlgo(BetterEnum):
    """
    The enum of Compress algorithms.

    .. note::
        LZ4 requires additional package, to install it use pip install "jina[lz4]"

    .. seealso::

        https://docs.python.org/3/library/archiving.html
    """

    NONE = 0
    LZ4 = 1
    ZLIB = 2
    GZIP = 3
    BZ2 = 4
    LZMA = 5


class OnErrorStrategy(BetterEnum):
    """
    The level of error handling.

    .. warning::
        In theory, all methods below do not 100% guarantee the success
        execution on the sequel flow. If something is wrong in the upstream,
        it is hard to CARRY this exception and moving forward without ANY
        side-effect.
    """

    IGNORE = (
        0  #: Ignore it, keep running all Drivers & Executors logics in the sequel flow
    )
    SKIP_HANDLE = 1  #: Skip all Executors in the sequel, only `pre_hook` and `post_hook` are called
    THROW_EARLY = 2  #: Immediately throw the exception, the sequel flow will not be running at all


class FlowInspectType(BetterEnum):
    """Inspect strategy in the flow."""

    HANG = 0  # keep them hanging there
    REMOVE = 1  # remove them in the build
    COLLECT = 2  # spawn a new deployment and collect them before build

    @property
    def is_keep(self) -> bool:
        """
        Check if the target is inspected.

        :return: True if the target is inspected else False.
        """
        return self.value in {0, 2}


class DataInputType(BetterEnum):
    """Data input type in the request generator."""

    AUTO = 0  # auto inference the input type from data (!WARN: could be slow as it relies on try-execept)
    DOCUMENT = 1  # the input is a full document
    CONTENT = 2  # the input is just the content of the document
    DICT = 3  # the input is a dictionary representing a Document, needed while pydantic model not available


class RuntimeBackendType(BetterEnum):
    """Type of backend in runtime."""

    THREAD = 0
    PROCESS = 1


class RemoteWorkspaceState(str, Enum):
    """Enum representing state of remote workspace created by JinaD"""

    PENDING = 'PENDING'
    CREATING = 'CREATING'
    UPDATING = 'UPDATING'
    ACTIVE = 'ACTIVE'
    FAILED = 'FAILED'
    DELETING = 'DELETING'


def replace_enum_to_str(obj):
    """
    Transform BetterEnum type into string.

    :param obj: Target obj.
    :return: Transformed obj with string type values.
    """
    for k, v in obj.items():
        if isinstance(v, dict):
            obj[k] = replace_enum_to_str(v)
        elif isinstance(v, BetterEnum):
            obj[k] = str(v)
    return obj