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, EnumMeta, Flag, IntEnum
from typing import List, Union


[docs]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)
[docs] @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
[docs]class BetterEnum(IntEnum, metaclass=EnumType): """The base class of Enum used in Jina.""" def __str__(self): return self.to_string()
[docs] def to_string(self): """ Convert the Enum to string representation :return: the string representation of the enum """ return self.name
def __format__(self, format_spec): # noqa """ override format method for python 3.7 :parameter format_spec: format_spec :return: format using actual value type unless __str__ has been overridden. """ # credit python 3.9 : https://github.com/python/cpython/blob/612019e60e3a5340542122dabbc7ce5a27a8c635/Lib/enum.py#L755 # fix to enum BetterEnum not correctly formated str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__) if self._member_type_ is object or str_overridden: cls = str val = str(self) # mix-in branch else: cls = self._member_type_ val = self._value_ return cls.__format__(val, format_spec)
[docs] @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)
[docs]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
[docs]class LogVerbosity(BetterEnum): """Verbosity level of the logger.""" DEBUG = 10 INFO = 20 SUCCESS = 25 WARNING = 30 ERROR = 40 CRITICAL = 50
[docs]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
[docs]class GatewayProtocolType(BetterEnum): """ Gateway communication protocol """ GRPC = 0 HTTP = 1 WEBSOCKET = 2
[docs] @classmethod def from_string_list(cls, string_list: List[Union[str, 'GatewayProtocolType']]): """ Returns a list of Enums from a list of strings or enums :param string_list: list of strings or enums :return: a list of Enums """ return [cls.from_string(s) if isinstance(s, str) else s for s in string_list]
[docs]class PodRoleType(BetterEnum): """The enum of a Pod role.""" HEAD = 0 WORKER = 1 GATEWAY = 2
[docs]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}
[docs]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}
[docs]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
[docs]class WebsocketSubProtocols(str, Enum): """Subprotocol supported with Websocket Gateway""" JSON = 'json' BYTES = 'bytes'
[docs]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