jina-ai · hanxiao · Jun 10, 2021 · Jun 5, 2021 · Jun 7, 2021 · Jun 7, 2021
diff --git a/.github/2.0/cookbooks/Flow.md b/.github/2.0/cookbooks/Flow.md
@@ -717,7 +717,7 @@ c.post('/')
 ```
 
 ```text
-         Client@27219[S]:connected to the gateway at 192.168.1.14:12345!
+         GRPCClient@27219[S]:connected to the gateway at 192.168.1.14:12345!
   |█                   | 📃    100 ⏱️ 0.0s 🐎 26690.1/s      1   requests takes 0 seconds (0.00s)
 	✅ done in ⏱ 0 seconds 🐎 24854.8/s
 ```
@@ -834,17 +834,17 @@ When use `curl`, make sure to pass the `-N/--no-buffer` flag.
 
 ### Python Client with REST Request
 
-In some case when `restful=True`, you may still need a Python client to query the server for debugging. You can use `WebsocketClient` in this case.
+In some case when `restful=True`, you may still need a Python client to query the server for debugging. You can use `Client` with an additional parameter `restful` in this case.
 
 ```python
-from jina.clients import WebSocketClient
+from jina import Client
 
-c = WebSocketClient(host='192.168.1.14', port_expose=12345)
+c = Client(host='192.168.1.14', port_expose=12345, restful=True)
 c.post('/')
 ```
 
 ```text
-WebSocketClient@27622[S]:Connected to the gateway at 192.168.1.14:12345
+         WebSocketClient@27622[S]:Connected to the gateway at 192.168.1.14:12345
   |█                   | 📃    100 ⏱️ 0.0s 🐎 19476.6/s      1   requests takes 0 seconds (0.00s)
 	✅ done in ⏱ 0 seconds 🐎 18578.9/s
 ```

diff --git a/jina/clients/__init__.py b/jina/clients/__init__.py
@@ -1,28 +1,79 @@
 """Module wrapping the Client of Jina."""
+import argparse
+from typing import overload, Optional
+
 from .base import BaseClient, CallbackFnType, InputType, InputDeleteType
 from .helper import callback_exec
 from .mixin import PostMixin
 from .request import GeneratorSourceType
 from .websocket import WebSocketClientMixin
 
-__all__ = ['Client', 'WebSocketClient']
+__all__ = ['Client', 'GRPCClient', 'WebSocketClient']
+
+
+# overload_inject_start_client
+@overload
+def Client(
+    continue_on_error: Optional[bool] = False,
+    host: Optional[str] = '0.0.0.0',
+    port_expose: Optional[int] = None,
+    proxy: Optional[bool] = False,
+    request_size: Optional[int] = 100,
+    restful: Optional[bool] = False,
+    return_results: Optional[bool] = False,
+    show_progress: Optional[bool] = False,
+    **kwargs
+) -> 'BaseClient':
+    """Create a Client. Client is how user interact with Flow
+
+    :param continue_on_error: If set, a Request that causes error will be logged only without blocking the further requests.
+    :param host: The host address of the runtime, by default it is 0.0.0.0.
+    :param port_expose: The port of the host exposed to the public
+    :param proxy: If set, respect the http_proxy and https_proxy environment variables. otherwise, it will unset these proxy variables before start. gRPC seems to prefer no proxy
+    :param request_size: The number of Documents in each Request.
+    :param restful: If set, use RESTful interface instead of gRPC as the main interface
+    :param return_results: This feature is only used for AsyncClient.
+
+          If set, the results of all Requests will be returned as a list. This is useful when one wants
+          process Responses in bulk instead of using callback.
+    :param show_progress: If set, client will show a progress bar on receiving every request.
+    :return: the new Client object
+
+    .. # noqa: DAR202
+    .. # noqa: DAR101
+    .. # noqa: DAR003
+    """
+    # overload_inject_end_client
+
 
+def Client(args: Optional['argparse.Namespace'] = None, **kwargs) -> 'BaseClient':
+    """Jina Python client.
 
-class Client(PostMixin, BaseClient):
+    :param args: Namespace args.
+    :param kwargs: Additional arguments.
+    :return: An instance of :class:`GRPCClient` or :class:`WebSocketClient`.
+    """
+    is_restful = kwargs.get('restful', False)
+    if is_restful:
+        return WebSocketClient(args, **kwargs)
+    else:
+        return GRPCClient(args, **kwargs)
+
+
+class GRPCClient(PostMixin, BaseClient):
     """A simple Python client for connecting to the gRPC gateway.
 
     It manages the asyncio event loop internally, so all interfaces are synchronous from the outside.
     """
 
     @property
-    def client(self) -> 'Client':
+    def client(self) -> 'GRPCClient':
         """Return the client object itself
-
         .. # noqa: DAR201"""
         return self
 
 
-class WebSocketClient(Client, WebSocketClientMixin):
+class WebSocketClient(GRPCClient, WebSocketClientMixin):
     """A Python Client to stream requests from a Flow with a REST Gateway.
 
     :class:`WebSocketClient` shares the same interface as :class:`Client` and provides methods like

diff --git a/jina/clients/mixin.py b/jina/clients/mixin.py
@@ -7,7 +7,7 @@
 
 
 class PostMixin:
-    """The Post Mixin class for Client and Flow """
+    """The Post Mixin class for Client and Flow"""
 
     def post(
         self,
@@ -63,7 +63,7 @@ async def _get_results(*args, **kwargs):
 
 
 class AsyncPostMixin:
-    """The Async Post Mixin class for AsyncClient and AsyncFlow """
+    """The Async Post Mixin class for AsyncClient and AsyncFlow"""
 
     async def post(
         self,

diff --git a/jina/flow/base.py b/jina/flow/base.py
@@ -12,7 +12,8 @@
 
 from .builder import build_required, _build_flow, _hanging_pods
 from .. import __default_host__
-from ..clients import Client, WebSocketClient
+from ..clients.base import BaseClient
+from ..clients import GRPCClient, WebSocketClient
 from ..enums import FlowBuildLevel, PodRoleType, FlowInspectType
 from ..excepts import FlowTopologyError, FlowMissingPodError
 from ..helper import (
@@ -68,7 +69,9 @@ class BaseFlow(JAMLCompatible, ExitStack, metaclass=FlowType):
     :param env: environment variables shared by all Pods
     """
 
-    _cls_client = Client  #: the type of the Client, can be changed to other class
+    _cls_client = (
+        GRPCClient  #: the type of the GRPCClient, can be changed to other class
+    )
 
     # overload_inject_start_flow
     @overload
@@ -211,7 +214,7 @@ def _add_gateway(self, needs, **kwargs):
                 name=pod_name,
                 ctrl_with_ipc=True,  # otherwise ctrl port would be conflicted
                 runtime_cls='GRPCRuntime'
-                if self._cls_client == Client
+                if self._cls_client == GRPCClient
                 else 'RESTRuntime',
                 pod_role=PodRoleType.GATEWAY,
                 identity=self.args.identity,
@@ -800,8 +803,8 @@ def __eq__(self, other: 'BaseFlow') -> bool:
 
     @property
     @build_required(FlowBuildLevel.GRAPH)
-    def client(self) -> 'Client':
-        """Return a :class:`Client` object attach to this Flow.
+    def client(self) -> 'BaseClient':
+        """Return a :class:`BaseClient` object attach to this Flow.
 
         .. # noqa: DAR201"""
         kwargs = {}
@@ -816,7 +819,6 @@ def client(self) -> 'Client':
         # show progress when client is used inside the flow, for better log readability
         if 'show_progress' not in kwargs:
             args.show_progress = True
-
         return self._cls_client(args)
 
     @property
@@ -1093,7 +1095,7 @@ def _show_success_message(self):
         self.logger.info('\n' + '\n'.join(address_table))
 
     def block(self):
-        """Block the process until user hits KeyboardInterrupt """
+        """Block the process until user hits KeyboardInterrupt"""
         try:
             threading.Event().wait()
         except KeyboardInterrupt:
@@ -1111,7 +1113,7 @@ def use_grpc_gateway(self, port: Optional[int] = None):
 
     def _switch_gateway(self, gateway: str, port: int):
         restful = gateway == 'RESTRuntime'
-        client = WebSocketClient if gateway == 'RESTRuntime' else Client
+        client = WebSocketClient if gateway == 'RESTRuntime' else GRPCClient
 
         # globally register this at Flow level
         self._cls_client = client

diff --git a/jina/parsers/client.py b/jina/parsers/client.py
@@ -42,3 +42,10 @@ def mixin_client_cli_parser(parser):
 process Responses in bulk instead of using callback. 
                     ''',
     )
+
+    gp.add_argument(
+        '--restful',
+        action='store_true',
+        default=False,
+        help='If set, use RESTful interface instead of gRPC as the main interface',
+    )
diff --git a/scripts/inject-cli-as-overload.py b/scripts/inject-cli-as-overload.py
@@ -57,6 +57,7 @@ def _cli_to_schema(
         return_type='BaseFlow',
         filepath='../jina/flow/base.py',
         overload_fn='add',
+        class_method=True,  # if it is a method inside class.
     ),
     dict(
         cli_entrypoint='flow',
@@ -65,6 +66,16 @@ def _cli_to_schema(
         return_type=None,
         filepath='../jina/flow/base.py',
         overload_fn='__init__',
+        class_method=True,
+    ),
+    dict(
+        cli_entrypoint='client',
+        doc_str_title='Create a Client. Client is how user interact with Flow',
+        doc_str_return='the new Client object',
+        return_type='BaseClient',
+        filepath='../jina/clients/__init__.py',
+        overload_fn='Client',
+        class_method=False,
     ),
 ]
 
@@ -76,37 +87,66 @@ def fill_overload(
     return_type,
     filepath,
     overload_fn,
+    class_method,
     indent=' ' * 4,
 ):
     a = _cli_to_schema(api_to_dict(), cli_entrypoint)
-    cli_args = [
-        f'{indent}{indent}{k[0]}: Optional[{k[1]["type"]}] = {k[1]["default_literal"]}'
-        for k in a
-    ]
-    args_str = ', \n'.join(cli_args + [f'{indent}{indent}**kwargs'])
-    signature_str = f'def {overload_fn}(\n{indent}{indent}self,\n{args_str})'
-    if return_type:
-        signature_str += f' -> \'{return_type}\':'
-        return_str = f'\n{indent}{indent}:return: {doc_str_return}'
+    if class_method:
+        cli_args = [
+            f'{indent}{indent}{k[0]}: Optional[{k[1]["type"]}] = {k[1]["default_literal"]}'
+            for k in a
+        ]
+        args_str = ', \n'.join(cli_args + [f'{indent}{indent}**kwargs'])
+        signature_str = f'def {overload_fn}(\n{indent}{indent}self,\n{args_str})'
+        if return_type:
+            signature_str += f' -> \'{return_type}\':'
+            return_str = f'\n{indent}{indent}:return: {doc_str_return}'
+        else:
+            signature_str += ':'
+            return_str = ''
+    else:
+        cli_args = [
+            f'{indent}{k[0]}: Optional[{k[1]["type"]}] = {k[1]["default_literal"]}'
+            for k in a
+        ]
+        args_str = ', \n'.join(cli_args + [f'{indent}**kwargs'])
+        signature_str = f'def {overload_fn}(\n{args_str})'
+        if return_type:
+            signature_str += f' -> \'{return_type}\':'
+            return_str = f'\n{indent}:return: {doc_str_return}'
+        else:
+            signature_str += ':'
+            return_str = ''
+    if class_method:
+        doc_str = '\n'.join(
+            f'{indent}{indent}:param {k[0]}: {k[1]["description"]}' for k in a
+        )
+        noqa_str = '\n'.join(
+            f'{indent}{indent}.. # noqa: DAR{j}' for j in ['202', '101', '003']
+        )
+    else:
+        doc_str = '\n'.join(f'{indent}:param {k[0]}: {k[1]["description"]}' for k in a)
+        noqa_str = '\n'.join(
+            f'{indent}.. # noqa: DAR{j}' for j in ['202', '101', '003']
+        )
+    if class_method:
+        final_str = f'@overload\n{indent}{signature_str}\n{indent}{indent}"""{doc_str_title}\n\n{doc_str}{return_str}\n\n{noqa_str}\n{indent}{indent}"""'
+        final_code = re.sub(
+            rf'(# overload_inject_start_{cli_entrypoint}).*(# overload_inject_end_{cli_entrypoint})',
+            f'\\1\n{indent}{final_str}\n{indent}\\2',
+            open(filepath).read(),
+            0,
+            re.DOTALL,
+        )
     else:
-        signature_str += ':'
-        return_str = ''
-    doc_str = '\n'.join(
-        f'{indent}{indent}:param {k[0]}: {k[1]["description"]}' for k in a
-    )
-    noqa_str = '\n'.join(
-        f'{indent}{indent}.. # noqa: DAR{j}' for j in ['202', '101', '003']
-    )
-
-    final_str = f'@overload\n{indent}{signature_str}\n{indent}{indent}"""{doc_str_title}\n\n{doc_str}{return_str}\n\n{noqa_str}\n{indent}{indent}"""'
-
-    final_code = re.sub(
-        rf'(# overload_inject_start_{cli_entrypoint}).*(# overload_inject_end_{cli_entrypoint})',
-        f'\\1\n{indent}{final_str}\n{indent}\\2',
-        open(filepath).read(),
-        0,
-        re.DOTALL,
-    )
+        final_str = f'@overload\n{signature_str}\n{indent}"""{doc_str_title}\n\n{doc_str}{return_str}\n\n{noqa_str}\n{indent}"""'
+        final_code = re.sub(
+            rf'(# overload_inject_start_{cli_entrypoint}).*(# overload_inject_end_{cli_entrypoint})',
+            f'\\1\n{final_str}\n{indent}\\2',
+            open(filepath).read(),
+            0,
+            re.DOTALL,
+        )
 
     with open(filepath, 'w') as fp:
         fp.write(final_code)

diff --git a/tests/distributed/test_join_local_from_remote/test_integration.py b/tests/distributed/test_join_local_from_remote/test_integration.py
@@ -3,7 +3,7 @@
 import pytest
 
 from jina import Document
-from jina.clients import Client
+from jina.clients import Client, GRPCClient
 from jina.parsers import set_client_cli_parser
 from tests import validate_callback
 from ..helpers import create_flow, assert_request
@@ -23,16 +23,26 @@ def doc_to_index():
 
 @pytest.fixture
 def client():
+    return Client(host='localhost', port_expose=45678)
+
+
+@pytest.fixture
+def grpc_client():
     args = set_client_cli_parser().parse_args(
         ['--host', 'localhost', '--port-expose', '45678']
     )
 
-    return Client(args)
    return GRPCClient(args)
+
+
+@pytest.fixture(params=['client', 'grpc_client'])
+def client_instance(request):
+    return request.getfixturevalue(request.param)
 
 
 @pytest.mark.timeout(360)
 @pytest.mark.parametrize('docker_compose', [compose_yml], indirect=['docker_compose'])
-def test_flow(docker_compose, doc_to_index, client, mocker):
+def test_flow(docker_compose, doc_to_index, client_instance, mocker):
     def validate_resp(resp):
         assert len(resp.data.docs) == 2
         assert resp.data.docs[0].text == 'test'
@@ -41,7 +51,7 @@ def validate_resp(resp):
     mock = mocker.Mock()
     flow_id = create_flow(flow_yaml=flow_yaml, pod_dir=os.path.join(cur_dir, 'pods'))
 
-    client.search(inputs=[doc_to_index], on_done=mock)
+    client_instance.search(inputs=[doc_to_index], on_done=mock)
 
     assert_request(method='get', url=f'http://localhost:8000/flows/{flow_id}')