From d244f5dd4f6472d197523a32331132cacb3ef16c Mon Sep 17 00:00:00 2001 From: Valay Dave Date: Tue, 15 Feb 2022 08:48:14 -0800 Subject: [PATCH] Adding export of graph to JSON for cli --- metaflow/cli.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/metaflow/cli.py b/metaflow/cli.py index 2674757aa82..1e5cde43e8f 100644 --- a/metaflow/cli.py +++ b/metaflow/cli.py @@ -174,10 +174,22 @@ def help(ctx): @cli.command(help="Output internal state of the flow graph.") +@click.option("--json", is_flag=True, help="Output the flow graph in JSON format.") @click.pass_obj -def output_raw(obj): - echo("Internal representation of the flow:", fg="magenta", bold=False) - echo_always(str(obj.graph), err=False) +def output_raw(obj, json): + if json: + import json as _json + + _msg = "Internal representation of the flow in JSON format:" + _graph_dict, _graph_struct = obj.graph.output_steps() + _graph = _json.dumps( + dict(graph=_graph_dict, graph_structure=_graph_struct), indent=4 + ) + else: + _graph = str(obj.graph) + _msg = "Internal representation of the flow:" + echo(_msg, fg="magenta", bold=False) + echo_always(_graph, err=False) @cli.command(help="Visualize the flow with Graphviz.")