add simple script generating graphviz from a balance

1 files changed, 39 insertions, 0 deletions

diff --git a/plot.py b/plot.py
new file mode 100755
index 0000000..b5c6355
--- /dev/null
+++ b/plot.py
@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+"""Plot a balance"""
+Balance = dict[str, dict[str, int]]
+
+
+def plot(balance: Balance, exclude=None) -> str:
+    """Plot a given balance"""
+
+    g = 'digraph balance {\n'
+    for u, b in balance.items():
+        if exclude and u in exclude:
+            continue
+
+        for other, val in b.items():
+            if exclude and other in exclude:
+                continue
+
+            if val > 0:
+                g += f'{u} -> {other} [label={val/100:.2f}];\n'
+
+    return g + '}'
+
+
+def main():
+    """Entry point listing minimize steps of a particular state file"""
+    import json
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('state_file', default='state.json', nargs='?')
+    parser.add_argument('-x', '--exclude', nargs='*')
+    args = parser.parse_args()
+    with open(args.state_file, 'r', encoding='utf-8') as state_file:
+        balance = json.load(state_file)['balance']
+    print(plot(balance, exclude=args.exclude))
+
+
+if __name__ == '__main__':
+    main()