From d3d0713de587c97c3c23200a15215eb0807ff87e Mon Sep 17 00:00:00 2001
From: Mark Haines <mark.haines@matrix.org>
Date: Thu, 15 Jan 2015 17:06:12 +0000
Subject: Move experiments, graph and cmdclient into contrib

---
 contrib/graph/graph.py | 151 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 151 insertions(+)
 create mode 100644 contrib/graph/graph.py

(limited to 'contrib/graph/graph.py')

diff --git a/contrib/graph/graph.py b/contrib/graph/graph.py
new file mode 100644
index 0000000000..b2acadcf5e
--- /dev/null
+++ b/contrib/graph/graph.py
@@ -0,0 +1,151 @@
+# Copyright 2014 OpenMarket Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import sqlite3
+import pydot
+import cgi
+import json
+import datetime
+import argparse
+import urllib2
+
+
+def make_name(pdu_id, origin):
+    return "%s@%s" % (pdu_id, origin)
+
+
+def make_graph(pdus, room, filename_prefix):
+    pdu_map = {}
+    node_map = {}
+
+    origins = set()
+    colors = set(("red", "green", "blue", "yellow", "purple"))
+
+    for pdu in pdus:
+        origins.add(pdu.get("origin"))
+
+    color_map = {color: color for color in colors if color in origins}
+    colors -= set(color_map.values())
+
+    color_map[None] = "black"
+
+    for o in origins:
+        if o in color_map:
+            continue
+        try:
+            c = colors.pop()
+            color_map[o] = c
+        except:
+            print "Run out of colours!"
+            color_map[o] = "black"
+
+    graph = pydot.Dot(graph_name="Test")
+
+    for pdu in pdus:
+        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')
+
+        label = (
+            "<"
+            "<b>%(name)s </b><br/>"
+            "Type: <b>%(type)s </b><br/>"
+            "State key: <b>%(state_key)s </b><br/>"
+            "Content: <b>%(content)s </b><br/>"
+            "Time: <b>%(time)s </b><br/>"
+            "Depth: <b>%(depth)s </b><br/>"
+            ">"
+        ) % {
+            "name": name,
+            "type": pdu.get("pdu_type"),
+            "state_key": pdu.get("state_key"),
+            "content": cgi.escape(json.dumps(pdu.get("content")), quote=True),
+            "time": t,
+            "depth": pdu.get("depth"),
+        }
+
+        node = pydot.Node(
+            name=name,
+            label=label,
+            color=color_map[pdu.get("origin")]
+        )
+        node_map[name] = node
+        graph.add_node(node)
+
+    for pdu in pdus:
+        start_name = make_name(pdu.get("pdu_id"), pdu.get("origin"))
+        for i, o in pdu.get("prev_pdus", []):
+            end_name = make_name(i, o)
+
+            if end_name not in node_map:
+                print "%s not in nodes" % end_name
+                continue
+
+            edge = pydot.Edge(node_map[start_name], node_map[end_name])
+            graph.add_edge(edge)
+
+        # Add prev_state edges, if they exist
+        if pdu.get("prev_state_id") and pdu.get("prev_state_origin"):
+            prev_state_name = make_name(
+                pdu.get("prev_state_id"), pdu.get("prev_state_origin")
+            )
+
+            if prev_state_name in node_map:
+                state_edge = pydot.Edge(
+                    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')
+
+
+def get_pdus(host, room):
+    transaction = json.loads(
+        urllib2.urlopen(
+            "http://%s/_matrix/federation/v1/context/%s/" % (host, room)
+        ).read()
+    )
+
+    return transaction["pdus"]
+
+
+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."
+    )
+    parser.add_argument(
+        "-p", "--prefix", dest="prefix",
+        help="String to prefix output files with"
+    )
+    parser.add_argument('host')
+    parser.add_argument('room')
+
+    args = parser.parse_args()
+
+    host = args.host
+    room = args.room
+    prefix = args.prefix if args.prefix else "%s_graph" % (room)
+
+    pdus = get_pdus(host, room)
+
+    make_graph(pdus, room, prefix)
-- 
cgit 1.4.1