Fix-up the contrib/graph scripts. (#13013)

* Clarifies comments and documentation.
* Adds type-hints.
* Fixes Python 3 compatibility (and runs pyupgrade).
* Updates for changes in Synapse internals.

3 files changed, 69 insertions, 43 deletions

diff --git a/contrib/graph/graph.py b/contrib/graph/graph.py
index fdbac087bd..3c4f47dbd2 100644
--- a/contrib/graph/graph.py
+++ b/contrib/graph/graph.py
@@ -1,11 +1,3 @@
-import argparse
-import cgi
-import datetime
-import json
-
-import pydot
-import urllib2
-
 # Copyright 2014-2016 OpenMarket Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,12 +12,25 @@ import urllib2
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import argparse
+import cgi
+import datetime
+import json
+import urllib.request
+from typing import List
+
+import pydot
+
 
-def make_name(pdu_id, origin):
-    return "%s@%s" % (pdu_id, origin)
+def make_name(pdu_id: str, origin: str) -> str:
+    return f"{pdu_id}@{origin}"
 
 
-def make_graph(pdus, room, filename_prefix):
+def make_graph(pdus: List[dict], filename_prefix: str) -> None:
+    """
+    Generate a dot and SVG file for a graph of events in the room based on the
+    topological ordering by querying a homeserver.
+    """
     pdu_map = {}
     node_map = {}
 
@@ -111,10 +116,10 @@ def make_graph(pdus, room, filename_prefix):
     graph.write_svg("%s.svg" % filename_prefix, prog="dot")
 
 
-def get_pdus(host, room):
+def get_pdus(host: str, room: str) -> List[dict]:
     transaction = json.loads(
-        urllib2.urlopen(
-            "http://%s/_matrix/federation/v1/context/%s/" % (host, room)
+        urllib.request.urlopen(
+            f"http://{host}/_matrix/federation/v1/context/{room}/"
         ).read()
     )
 
@@ -141,4 +146,4 @@ if __name__ == "__main__":
 
     pdus = get_pdus(host, room)
 
-    make_graph(pdus, room, prefix)
+    make_graph(pdus, prefix)
diff --git a/contrib/graph/graph2.py b/contrib/graph/graph2.py
index 0980231e4a..b46094ce0a 100644
--- a/contrib/graph/graph2.py
+++ b/contrib/graph/graph2.py
@@ -14,22 +14,31 @@
 
 
 import argparse
-import cgi
 import datetime
+import html
 import json
 import sqlite3
 
 import pydot
 
-from synapse.events import FrozenEvent
+from synapse.api.room_versions import KNOWN_ROOM_VERSIONS
+from synapse.events import make_event_from_dict
 from synapse.util.frozenutils import unfreeze
 
 
-def make_graph(db_name, room_id, file_prefix, limit):
+def make_graph(db_name: str, room_id: str, file_prefix: str, limit: int) -> None:
+    """
+    Generate a dot and SVG file for a graph of events in the room based on the
+    topological ordering by reading from a Synapse SQLite database.
+    """
     conn = sqlite3.connect(db_name)
 
+    sql = "SELECT room_version FROM rooms WHERE room_id = ?"
+    c = conn.execute(sql, (room_id,))
+    room_version = KNOWN_ROOM_VERSIONS[c.fetchone()[0]]
+
     sql = (
-        "SELECT json FROM event_json as j "
+        "SELECT json, internal_metadata FROM event_json as j "
         "INNER JOIN events as e ON e.event_id = j.event_id "
         "WHERE j.room_id = ?"
     )
@@ -43,7 +52,10 @@ def make_graph(db_name, room_id, file_prefix, limit):
 
     c = conn.execute(sql, args)
 
-    events = [FrozenEvent(json.loads(e[0])) for e in c.fetchall()]
+    events = [
+        make_event_from_dict(json.loads(e[0]), room_version, json.loads(e[1]))
+        for e in c.fetchall()
+    ]
 
     events.sort(key=lambda e: e.depth)
 
@@ -84,7 +96,7 @@ def make_graph(db_name, room_id, file_prefix, limit):
             "name": event.event_id,
             "type": event.type,
             "state_key": event.get("state_key", None),
-            "content": cgi.escape(content, quote=True),
+            "content": html.escape(content, quote=True),
             "time": t,
             "depth": event.depth,
             "state_group": state_group,
@@ -96,11 +108,11 @@ def make_graph(db_name, room_id, file_prefix, limit):
         graph.add_node(node)
 
     for event in events:
-        for prev_id, _ in event.prev_events:
+        for prev_id in event.prev_event_ids():
             try:
                 end_node = node_map[prev_id]
             except Exception:
-                end_node = pydot.Node(name=prev_id, label="<<b>%s</b>>" % (prev_id,))
+                end_node = pydot.Node(name=prev_id, label=f"<<b>{prev_id}</b>>")
 
                 node_map[prev_id] = end_node
                 graph.add_node(end_node)
@@ -112,7 +124,7 @@ def make_graph(db_name, room_id, file_prefix, limit):
         if len(event_ids) <= 1:
             continue
 
-        cluster = pydot.Cluster(str(group), label="<State Group: %s>" % (str(group),))
+        cluster = pydot.Cluster(str(group), label=f"<State Group: {str(group)}>")
 
         for event_id in event_ids:
             cluster.add_node(node_map[event_id])
@@ -126,7 +138,7 @@ def make_graph(db_name, room_id, file_prefix, limit):
 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"
+        "to the given Synapse SQLite file to get the list of PDUs. \n"
         "Requires pydot."
     )
     parser.add_argument(
diff --git a/contrib/graph/graph3.py b/contrib/graph/graph3.py
index dd0c19368b..a28a1594c7 100644
--- a/contrib/graph/graph3.py
+++ b/contrib/graph/graph3.py
@@ -1,13 +1,3 @@
-import argparse
-import cgi
-import datetime
-
-import pydot
-import simplejson as json
-
-from synapse.events import FrozenEvent
-from synapse.util.frozenutils import unfreeze
-
 # Copyright 2016 OpenMarket Ltd
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -22,15 +12,35 @@ from synapse.util.frozenutils import unfreeze
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import argparse
+import datetime
+import html
+import json
+
+import pydot
 
-def make_graph(file_name, room_id, file_prefix, limit):
+from synapse.api.room_versions import KNOWN_ROOM_VERSIONS
+from synapse.events import make_event_from_dict
+from synapse.util.frozenutils import unfreeze
+
+
+def make_graph(file_name: str, file_prefix: str, limit: int) -> None:
+    """
+    Generate a dot and SVG file for a graph of events in the room based on the
+    topological ordering by reading line-delimited JSON from a file.
+    """
     print("Reading lines")
     with open(file_name) as f:
         lines = f.readlines()
 
     print("Read lines")
 
-    events = [FrozenEvent(json.loads(line)) for line in lines]
+    # Figure out the room version, assume the first line is the create event.
+    room_version = KNOWN_ROOM_VERSIONS[
+        json.loads(lines[0]).get("content", {}).get("room_version")
+    ]
+
+    events = [make_event_from_dict(json.loads(line), room_version) for line in lines]
 
     print("Loaded events.")
 
@@ -66,8 +76,8 @@ def make_graph(file_name, room_id, file_prefix, limit):
             content.append(
                 "<b>%s</b>: %s,"
                 % (
-                    cgi.escape(key, quote=True).encode("ascii", "xmlcharrefreplace"),
-                    cgi.escape(value, quote=True).encode("ascii", "xmlcharrefreplace"),
+                    html.escape(key, quote=True).encode("ascii", "xmlcharrefreplace"),
+                    html.escape(value, quote=True).encode("ascii", "xmlcharrefreplace"),
                 )
             )
 
@@ -101,11 +111,11 @@ def make_graph(file_name, room_id, file_prefix, limit):
     print("Created Nodes")
 
     for event in events:
-        for prev_id, _ in event.prev_events:
+        for prev_id in event.prev_event_ids():
             try:
                 end_node = node_map[prev_id]
             except Exception:
-                end_node = pydot.Node(name=prev_id, label="<<b>%s</b>>" % (prev_id,))
+                end_node = pydot.Node(name=prev_id, label=f"<<b>{prev_id}</b>>")
 
                 node_map[prev_id] = end_node
                 graph.add_node(end_node)
@@ -139,8 +149,7 @@ if __name__ == "__main__":
     )
     parser.add_argument("-l", "--limit", help="Only retrieve the last N events.")
     parser.add_argument("event_file")
-    parser.add_argument("room")
 
     args = parser.parse_args()
 
-    make_graph(args.event_file, args.room, args.prefix, args.limit)
+    make_graph(args.event_file, args.prefix, args.limit)