diff --git a/gnes/flow/__init__.py b/gnes/flow/__init__.py index 504a29d9..cfda7d96 100644 --- a/gnes/flow/__init__.py +++ b/gnes/flow/__init__.py @@ -139,6 +139,12 @@ def from_yaml(orchestration: str) -> 'Flow': @_build_level(BuildLevel.GRAPH) def to_mermaid(self, left_right: bool = True): + """ + Output the mermaid graph for visualization + + :param left_right: render the flow in left-to-right manner, otherwise top-down manner. + :return: + """ mermaid_graph = OrderedDict() for k in self._service_nodes.keys(): mermaid_graph[k] = [] @@ -166,10 +172,28 @@ def to_mermaid(self, left_right: bool = True): ['graph %s' % ('LR' if left_right else 'TD')] + [ss for s in mermaid_graph.values() for ss in s] + style + class_def) - self.logger.info( - 'copy-paste the output and visualize it with: https://mermaidjs.github.io/mermaid-live-editor/') return mermaid_str + @_build_level(BuildLevel.GRAPH) + def to_jpg(self, path: str = 'flow.jpg', left_right: bool = True): + """ + Rendering the current flow as a jpg image, this will call :py:meth:`to_mermaid` and it needs internet connection + + :param path: the file path of the image + :param left_right: render the flow in left-to-right manner, otherwise top-down manner. + :return: + """ + import base64 + from urllib.request import Request, urlopen + mermaid_str = self.to_mermaid(left_right) + encoded_str = base64.b64encode(bytes(mermaid_str, 'utf-8')).decode('utf-8') + print('https://mermaidjs.github.io/mermaid-live-editor/#/view/%s' % encoded_str) + self.logger.info('saving jpg...') + req = Request('https://mermaid.ink/img/%s' % encoded_str, headers={'User-Agent': 'Mozilla/5.0'}) + with open(path, 'wb') as fp: + fp.write(urlopen(req).read()) + self.logger.info('done') + def train(self, bytes_gen: Iterator[bytes] = None, **kwargs): """Do training on the current flow diff --git a/tests/test_gnes_flow.py b/tests/test_gnes_flow.py index 10b2429b..b8bf4b5e 100644 --- a/tests/test_gnes_flow.py +++ b/tests/test_gnes_flow.py @@ -113,6 +113,7 @@ def test_flow5(self): .build(backend=None)) print(f._service_edges) print(f.to_mermaid()) + f.to_jpg() def _test_index_flow(self, backend): for k in [self.indexer1_bin, self.indexer2_bin, self.encoder_bin]: