From 9b4a5c02dcd022cbfe3ec89b405900da20f1a885 Mon Sep 17 00:00:00 2001 From: theUnkownName Date: Tue, 6 Oct 2020 10:10:01 +0200 Subject: [PATCH] fix: update docstring and remove return url --- jina/flow/__init__.py | 23 ++++++++++++++-------- tests/unit/flow/test_flow_visualization.py | 10 ---------- 2 files changed, 15 insertions(+), 18 deletions(-) diff --git a/jina/flow/__init__.py b/jina/flow/__init__.py index 3401f1e6c3c10..a89ffc91b8966 100644 --- a/jina/flow/__init__.py +++ b/jina/flow/__init__.py @@ -170,7 +170,6 @@ def __init__(self, args: 'argparse.Namespace' = None, **kwargs): self._build_level = FlowBuildLevel.EMPTY self._pod_name_counter = 0 self._last_changed_pod = ['gateway'] #: default first pod is gateway, will add when build() - self._url = '' self._update_args(args, **kwargs) def _update_args(self, args, **kwargs): @@ -790,14 +789,22 @@ def my_reader(): """ self._get_client(**kwargs).search(input_fn, output_fn, **kwargs) - def plot(self, output='', copy_flow: bool = True) -> 'Flow': + def plot(self, output: str=None, copy_flow: bool = False) -> 'Flow': """ - Output the mermaid graph for visualization - :return: the flow - """ + Creates a mermaid graph for the Flow. + If a file name is provided it will create a jpg image with that name, + otherwise it will display the URL for mermaid + Example, + .. highlight:: python + .. code-block:: python + flow = Flow().add(name='pod_a').plot('flow_test.jpg') + + :param output: a filename specifying the name of the image to be created + :return: the flow + """ + op_flow = copy.deepcopy(self) if copy_flow else self url = '' - mermaid_graph = ['graph TD'] for node, v in self._pod_nodes.items(): for need in sorted(v.needs): @@ -806,7 +813,7 @@ def plot(self, output='', copy_flow: bool = True) -> 'Flow': if output: self._mermaidstr_to_jpg(mermaid_str, output) else: - self._url = self._mermaidstr_to_url(mermaid_str) + url = self._mermaidstr_to_url(mermaid_str) return op_flow @@ -838,7 +845,7 @@ def _mermaidstr_to_jpg(self, mermaid_str, output) -> None: req = Request('https://mermaid.ink/img/%s' % encoded_str, headers={'User-Agent': 'Mozilla/5.0'}) with open(output, 'wb') as fp: fp.write(urlopen(req).read()) - self.logger.info('done') + self.logger.info('dreturn 'https://mermaidjs.github.io/mermaid-live-editor/#/view/' + encoded_strone') def dry_run(self, **kwargs): """Send a DRYRUN request to this flow, passing through all pods in this flow, diff --git a/tests/unit/flow/test_flow_visualization.py b/tests/unit/flow/test_flow_visualization.py index f6070bfa210b3..2bb12dc0afd98 100644 --- a/tests/unit/flow/test_flow_visualization.py +++ b/tests/unit/flow/test_flow_visualization.py @@ -4,16 +4,6 @@ from jina.flow import Flow -def test_visualization_url(): - flow = (Flow().add(name='pod_a') - .add(name='pod_b', needs='gateway') - .join(needs=['pod_a', 'pod_b']).plot()) - - url_split = flow._url.split('view/') #check that has info after standard URL text - - assert url_split is not ' ' - - def test_visualization_with_yml_file_img(tmpdir): cur_dir = os.path.dirname(os.path.abspath(__file__)) tmpfile = os.path.join(tmpdir, 'flow_test.jpg')