feat(jaml): allow loading from flatten Flow yaml (#2216)

jina/jaml/__init__.py

@@ -3,7 +3,7 @@
 import tempfile
 import warnings
 from types import SimpleNamespace
-from typing import Dict, Any, Union, TextIO, Optional, List
+from typing import Dict, Any, Union, TextIO, Optional, List, Tuple
 
 import yaml
 from yaml.constructor import FullConstructor
@@ -108,16 +108,23 @@ def escape(value: str, include_unknown_tags: bool = True) -> str:
         return re.sub(r, r'jtype: \1', value)
 
     @staticmethod
-    def unescape(value: str, include_unknown_tags: bool = True) -> str:
+    def unescape(value: str,
+                 include_unknown_tags: bool = True,
+                 jtype_whitelist: Tuple[str, ...] = None,
+                 ) -> str:
         """
         Unescape the YAML content by replacing all ``jtype: `` to tags.
 
         :param value: the escaped YAML content
         :param include_unknown_tags: if to include unknown tags during unescaping
+        :param jtype_whitelist: the list of jtype to be unescaped
         :return: unescaped YAML
         """
         if include_unknown_tags:
             r = r'jtype: (\w+)\b'
+        elif jtype_whitelist:
+            r = '|'.join(jtype_whitelist)
+            r = rf'jtype: ({r})\b'
         else:
             r = '|'.join(JAML.registered_tags())
             r = rf'jtype: ({r})\b'
@@ -507,8 +514,14 @@ def load_config(
                     no_tag_yml,
                     extra_search_paths=(os.path.dirname(s_path),) if s_path else None,
                 )
-            # revert yaml's tag and load again, this time with substitution
-            tag_yml = JAML.unescape(JAML.dump(no_tag_yml))
+            from ..flow import BaseFlow
+            if issubclass(cls, BaseFlow):
+                tag_yml = JAML.unescape(JAML.dump(no_tag_yml),
+                                        include_unknown_tags=False,
+                                        jtype_whitelist=('Flow', 'AsyncFlow'))
+            else:
+                # revert yaml's tag and load again, this time with substitution
+                tag_yml = JAML.unescape(JAML.dump(no_tag_yml))
 
             # load into object, no more substitute
             return JAML.load(tag_yml, substitute=False)

jina/parsers/helloworld.py

@@ -188,14 +188,6 @@ def set_hw_chatbot_parser(parser=None):
         parser = set_base_parser()
 
     mixin_hw_base_parser(parser)
-    parser.add_argument(
-        '--uses',
-        type=str,
-        default=resource_filename(
-            'jina', '/'.join(('resources', 'chatbot', 'helloworld.flow.index.yml'))
-        ),
-        help='The yaml path of the index flow',
-    )
     parser.add_argument(
         '--index-data-url',
         type=str,

tests/unit/flow/test_flow.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+
 from jina import Flow, Document
 from jina.enums import SocketType, FlowBuildLevel
 from jina.excepts import RuntimeFailToStart
@@ -755,3 +756,17 @@ class CustomizedExecutor(BaseExecutor):
 
     with f:
         pass
+
+
+def test_flow_allinone_yaml():
+    from jina import Encoder
+    class CustomizedEncoder(Encoder):
+        pass
+
+    f = Flow.load_config(os.path.join(cur_dir, 'yaml/flow-allinone.yml'))
+    with f:
+        pass
+
+    f = Flow.load_config(os.path.join(cur_dir, 'yaml/flow-allinone-oldstyle.yml'))
+    with f:
+        pass

tests/unit/flow/yaml/flow-allinone-oldstyle.yml

@@ -0,0 +1,22 @@
+jtype: Flow
+version: '1'
+pods:
+  - uses:
+    !CustomizedEncoder
+    requests:
+      on:
+        IndexRequest:
+          - !FilterQL
+            with:
+              lookups: { modality__in: [ mode2 ] }
+              traversal_paths: [ 'c' ]
+          - !EncodeDriver
+            with:
+              method: encode
+              traversal_paths: [ 'c' ]
+  - uses: _index
+  - uses:
+    !BaseEncoder
+    metas:
+      name: test_indexer
+      workspace: ./indexed

tests/unit/flow/yaml/flow-allinone.yml

@@ -0,0 +1,22 @@
+jtype: Flow
+version: '1'
+pods:
+  - uses:
+    jtype: CustomizedEncoder
+    requests:
+      on:
+        IndexRequest:
+          - jtype: FilterQL
+            with:
+              lookups: { modality__in: [ mode2 ] }
+              traversal_paths: [ 'c' ]
+          - jtype: EncodeDriver
+            with:
+              method: encode
+              traversal_paths: [ 'c' ]
+  - uses: _index
+  - uses:
+    jtype: BaseEncoder
+    metas:
+      name: test_indexer
+      workspace: ./indexed