bricakeld · April 7, 2019 07:21 · Apr 7, 2019
diff --git a/tree_truncate.py b/tree_truncate.py
@@ -0,0 +1,62 @@
+import sys
+import os
+import re
+
+dir = os.path.dirname(__file__)
+
+def usage():
+    print('python tree_truncate.py <input dot file> <number of nodes> <optional: output dot file | default=out.dot>')
+
+def check():
+    if len(sys.argv) < 3 or len(sys.argv) > 4:
+        usage()
+        sys.exit()
+
+def truncate():
+    infile = os.path.abspath(sys.argv[1])
+    outfile = 'out.dot'
+    num_nodes = int(sys.argv[2])
+    print(infile)
+
+    if len(sys.argv) == 4:
+        outfile = os.path.abspath(sys.argv[3])
+
+    content = None
+    with open(infile) as f:
+        content = f.readlines()
+
+    content = [x.strip() for x in content]
+
+    if len(content) < (num_nodes * 2) + 4:
+        print('there are less nodes than requested, no need to truncate')
+        sys.exit()
+
+    outlines = []
+    outlines.extend(content[:4])
+
+    nums = [str(x) for x in range(num_nodes)]
+    matchers = ['^' + x + ' ->' for x in nums]
+    matchers.extend(['^' + x + ' \[' for x in nums])
+    matchers.extend([' -> ' + x + ' ' for x in nums])
+    print(matchers)
+
+    outlines.extend([s for s in content if any(bool(re.search(xs, s)) for xs in matchers)])
+
+    for l in outlines:
+        if re.match('^[\d]+ -> [\d]+ ;$', l):
+            temp = l.split()
+            number = temp[2]
+            if int(number) > num_nodes:
+                outlines.extend([s for s in content if bool(re.search('^' + number + ' \[', s))])
+
+    # print(outlines)
+    outlines.extend(content[-1])
+    with open(outfile, 'w', encoding='utf-8') as f:
+        f.writelines('%s\n' % s for s in outlines)
+
+def main():
+    check()
+    truncate()
+
+if __name__ == '__main__':
+    main()
No results found