summary refs log tree commit diff
path: root/contrib/graph/graph.py
diff options
context:
space:
mode:
Diffstat (limited to 'contrib/graph/graph.py')
-rw-r--r--contrib/graph/graph.py38
1 files changed, 17 insertions, 21 deletions
diff --git a/contrib/graph/graph.py b/contrib/graph/graph.py
index afd1d446b4..92736480eb 100644
--- a/contrib/graph/graph.py
+++ b/contrib/graph/graph.py
@@ -1,3 +1,5 @@
+from __future__ import print_function
+
 # Copyright 2014-2016 OpenMarket Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -48,7 +50,7 @@ def make_graph(pdus, room, filename_prefix):
             c = colors.pop()
             color_map[o] = c
         except:
-            print "Run out of colours!"
+            print("Run out of colours!")
             color_map[o] = "black"
 
     graph = pydot.Dot(graph_name="Test")
@@ -57,9 +59,9 @@ def make_graph(pdus, room, filename_prefix):
         name = make_name(pdu.get("pdu_id"), pdu.get("origin"))
         pdu_map[name] = pdu
 
-        t = datetime.datetime.fromtimestamp(
-            float(pdu["ts"]) / 1000
-        ).strftime('%Y-%m-%d %H:%M:%S,%f')
+        t = datetime.datetime.fromtimestamp(float(pdu["ts"]) / 1000).strftime(
+            "%Y-%m-%d %H:%M:%S,%f"
+        )
 
         label = (
             "<"
@@ -79,11 +81,7 @@ def make_graph(pdus, room, filename_prefix):
             "depth": pdu.get("depth"),
         }
 
-        node = pydot.Node(
-            name=name,
-            label=label,
-            color=color_map[pdu.get("origin")]
-        )
+        node = pydot.Node(name=name, label=label, color=color_map[pdu.get("origin")])
         node_map[name] = node
         graph.add_node(node)
 
@@ -93,7 +91,7 @@ def make_graph(pdus, room, filename_prefix):
             end_name = make_name(i, o)
 
             if end_name not in node_map:
-                print "%s not in nodes" % end_name
+                print("%s not in nodes" % end_name)
                 continue
 
             edge = pydot.Edge(node_map[start_name], node_map[end_name])
@@ -107,14 +105,13 @@ def make_graph(pdus, room, filename_prefix):
 
             if prev_state_name in node_map:
                 state_edge = pydot.Edge(
-                    node_map[start_name], node_map[prev_state_name],
-                    style='dotted'
+                    node_map[start_name], node_map[prev_state_name], style="dotted"
                 )
                 graph.add_edge(state_edge)
 
-    graph.write('%s.dot' % filename_prefix, format='raw', prog='dot')
-#    graph.write_png("%s.png" % filename_prefix, prog='dot')
-    graph.write_svg("%s.svg" % filename_prefix, prog='dot')
+    graph.write("%s.dot" % filename_prefix, format="raw", prog="dot")
+    #    graph.write_png("%s.png" % filename_prefix, prog='dot')
+    graph.write_svg("%s.svg" % filename_prefix, prog="dot")
 
 
 def get_pdus(host, room):
@@ -130,15 +127,14 @@ def get_pdus(host, room):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="Generate a PDU graph for a given room by talking "
-                    "to the given homeserver to get the list of PDUs. \n"
-                    "Requires pydot."
+        "to the given homeserver to get the list of PDUs. \n"
+        "Requires pydot."
     )
     parser.add_argument(
-        "-p", "--prefix", dest="prefix",
-        help="String to prefix output files with"
+        "-p", "--prefix", dest="prefix", help="String to prefix output files with"
     )
-    parser.add_argument('host')
-    parser.add_argument('room')
+    parser.add_argument("host")
+    parser.add_argument("room")
 
     args = parser.parse_args()