add more print format to print command line (#348)

* add more print format to print command line

* style

* dic

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,7 @@ Change Logs
 0.8.5
 +++++
 
+* :pr:`348`: add format dot, shape to command line print
 * :pr:`346`: fix patch for sdpa_mask_recent_torch even if it was removed in transformers>=5.0
 
 0.8.4

_unittests/ut_xrun_doc/test_command_lines_exe.py

@@ -18,7 +18,7 @@ def dummy_path(self):
         )
 
     def test_a_parser_print(self):
-        for fmt in ["raw", "text", "pretty", "printer"]:
+        for fmt in ["raw", "text", "pretty", "printer", "shape", "dot"]:
             with self.subTest(format=fmt):
                 st = StringIO()
                 with redirect_stdout(st):
@@ -199,7 +199,9 @@ def forward(self, x):
         with redirect_stdout(st):
             main(args)
         text = st.getvalue()
-        print(text)
+        if text:
+            # text is empty is dot is not installed
+            self.assertIn("converts into dot", text)
 
 
 if __name__ == "__main__":

onnx_diagnostic/_command_lines_parser.py

@@ -198,15 +198,19 @@ def get_parser_print() -> ArgumentParser:
     )
     parser.add_argument(
         "fmt",
-        choices=["pretty", "raw", "text", "printer"],
+        choices=["dot", "pretty", "printer", "raw", "shape", "text"],
         default="pretty",
         help=textwrap.dedent(
             """
             Prints out a model on the standard output.
-            raw     - just prints the model with print(...)
-            printer - onnx.printer.to_text(...)
+
+            dot     - converts the graph into dot
             pretty  - an improved rendering
+            printer - onnx.printer.to_text(...)
+            raw     - just prints the model with print(...)
+            shape   - prints every node node with input and output shapes
             text    - uses GraphRendering
+
             """.strip(
                 "\n"
             )
@@ -232,6 +236,14 @@ def _cmd_print(argv: List[Any]):
         from .helpers.graph_helper import GraphRendering
 
         print(GraphRendering(onx).text_rendering())
+    elif args.fmt == "shape":
+        from experimental_experiment.xbuilder import GraphBuilder
+
+        print(GraphBuilder(onx).pretty_text())
+    elif args.fmt == "dot":
+        from .helpers.dot_helper import to_dot
+
+        print(to_dot(onx))
     else:
         raise ValueError(f"Unexpected value fmt={args.fmt!r}")