Fix test stat coverage script (onnx#1948)

hariharans29 · Apr 18, 2019 · bff40c8 · bff40c8
1 parent aeb721a
commit bff40c8
Showing 1 changed file with 14 additions and 10 deletions.
diff --git a/onnx/backend/test/stat_coverage.py b/onnx/backend/test/stat_coverage.py
@@ -228,17 +228,21 @@ def main():
     docs_dir = os.path.join(base_dir, 'docs')
     schemas = defs.get_all_schemas()
 
-    ml = is_ml(schemas)
-    if ml:
-        fname = os.path.join(docs_dir, 'TestCoverage-ml.md')
-    else:
-        fname = os.path.join(docs_dir, 'TestCoverage.md')
-
+    has_ml = is_ml(schemas)
+    fname = os.path.join(docs_dir, 'TestCoverage.md')
     with io.open(fname, 'w+', newline='', encoding="utf-8") as f:  # type: ignore
-        gen_outlines(f, ml)
-        gen_node_test_coverage(schemas, f, ml)
-        gen_model_test_coverage(schemas, f, ml)
-        gen_overall_test_coverage(schemas, f, ml)
+        gen_outlines(f, False)
+        gen_node_test_coverage(schemas, f, False)
+        gen_model_test_coverage(schemas, f, False)
+        gen_overall_test_coverage(schemas, f, False)
+
+    if has_ml:
+        fname = os.path.join(docs_dir, 'TestCoverage-ml.md')
+        with io.open(fname, 'w+', newline='', encoding="utf-8") as f:  # type: ignore
+            gen_outlines(f, True)
+            gen_node_test_coverage(schemas, f, True)
+            gen_model_test_coverage(schemas, f, True)
+            gen_overall_test_coverage(schemas, f, True)
 
 
 if __name__ == '__main__':