Merge pull request #89 from nschloe/small-improv

some small improvements

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = tuna
-version = 0.5.0
+version = 0.5.1
 author = Nico Schlömer
 author_email = nico.schloemer@gmail.com
 description = Visualize Python performance profiles

test/test_tuna.py

@@ -26,19 +26,19 @@ def test_importprofile():
 """
 
     ref = {
-        "name": "main",
+        "text": ["main"],
         "color": 0,
         "children": [
             {
-                "name": "a",
+                "text": ["a"],
                 "value": 1e-06,
                 "color": 0,
                 "children": [
                     {
-                        "name": "b",
+                        "text": ["b"],
                         "value": 2e-06,
                         "color": 0,
-                        "children": [{"name": "c", "value": 3e-06, "color": 0}],
+                        "children": [{"text": ["c"], "value": 3e-06, "color": 0}],
                     }
                 ],
             }
@@ -66,23 +66,23 @@ def test_importprofile_multiprocessing():
 import time:       1 |    12 | a
 """
     ref = {
-        "name": "main",
+        "text": ["main"],
         "color": 0,
         "children": [
             {
-                "name": "a",
+                "text": ["a"],
                 "value": 1e-06,
                 "children": [
                     {
-                        "name": "b",
+                        "text": ["b"],
                         "value": 2e-06,
                         "children": [
                             {
-                                "name": "c",
+                                "text": ["c"],
                                 "value": 3e-06,
                                 "children": [
                                     {
-                                        "name": "e",
+                                        "text": ["e"],
                                         "value": 4.9999999999999996e-06,
                                         "color": 0,
                                     }
@@ -92,7 +92,7 @@ def test_importprofile_multiprocessing():
                         ],
                         "color": 0,
                     },
-                    {"name": "d", "value": 4e-06, "color": 0},
+                    {"text": ["d"], "value": 4e-06, "color": 0},
                 ],
                 "color": 0,
             }

tuna/_helpers.py

@@ -0,0 +1,2 @@
+class TunaError(Exception):
+    pass

tuna/_import_profile.py

@@ -0,0 +1,101 @@
+import logging
+
+from ._helpers import TunaError
+from .module_groups import built_in, built_in_deprecated
+
+
+def read_import_profile(filename):
+    # The import profile is of the form
+    # ```
+    # import time: self [us] | cumulative | imported package
+    # import time:       378 |        378 | zipimport
+    # import time:      1807 |       1807 | _frozen_importlib_external
+    # import time:       241 |        241 |     _codecs
+    # import time:      6743 |       6984 |   codecs
+    # import time:      1601 |       1601 |   encodings.aliases
+    # import time:     11988 |      20571 | encodings
+    # import time:       700 |        700 | encodings.utf_8
+    # import time:       535 |        535 | _signal
+    # import time:      1159 |       1159 | encodings.latin_1
+    # [...]
+    # ```
+    # The indentation in the last column signals parent-child relationships. In the
+    # above example, `encodings` is parent to `encodings.aliases` and `codecs` which in
+    # turn is parent to `_codecs`.
+    entries = []
+    with open(filename) as f:
+        # filtered iterator over lines prefixed with "import time: "
+        try:
+            # skip first line
+            next(f)
+        except UnicodeError:
+            raise TunaError()
+
+        for line in f:
+            if not line.startswith("import time: "):
+                logging.warning(f"Didn't recognize and skipped line `{line.rstrip()}`")
+                continue
+
+            line = line[len("import time: ") :].rstrip()
+
+            if line == "self [us] | cumulative | imported package":
+                continue
+            items = line.split(" | ")
+            assert len(items) == 3
+            self_time = int(items[0])
+            last = items[2]
+            name = last.lstrip()
+            num_leading_spaces = len(last) - len(name)
+            assert num_leading_spaces % 2 == 0
+            indentation_level = num_leading_spaces // 2 + 1
+            entries.append((name, indentation_level, self_time))
+
+    tree = _sort_into_tree(entries[::-1])
+
+    # go through the tree and add "color"
+    _add_color(tree, False)
+    return tree[0]
+
+
+def _add_color(tree, ancestor_is_built_in):
+    for item in tree:
+        module_name = item["text"][0].split(".")[0]
+        is_built_in = (
+            ancestor_is_built_in
+            or module_name in built_in
+            or module_name in built_in_deprecated
+        )
+        color = 1 if is_built_in else 0
+        if module_name in built_in_deprecated:
+            color = 2
+        item["color"] = color
+        if "children" in item:
+            _add_color(item["children"], is_built_in)
+
+
+def _sort_into_tree(lst):
+    main = {"text": ["main"], "color": 0, "children": []}
+
+    # keep a dictionary of the last entry of any given level
+    last = {}
+    last[0] = main
+
+    for entry in lst:
+        name, level, time = entry
+        # find the last entry with level-1
+        last[level - 1]["children"] += [
+            {"text": [name], "value": time * 1.0e-6, "children": []}
+        ]
+        last[level] = last[level - 1]["children"][-1]
+
+    _remove_empty_children(main)
+    return [main]
+
+
+def _remove_empty_children(tree):
+    if not tree["children"]:
+        del tree["children"]
+    else:
+        for k, child in enumerate(tree["children"]):
+            tree["children"][k] = _remove_empty_children(child)
+    return tree

tuna/_runtime_profile.py

@@ -0,0 +1,90 @@
+import pstats
+
+
+def read_runtime_profile(prof_filename):
+    stats = pstats.Stats(prof_filename)
+
+    # One way of picking the root nodes would be to search through stats.stats.items()
+    # and check which don't have parents. This, however, doesn't work if there are loops
+    # in the graph which happens, for example, if exec() is called somewhere in the
+    # program. For this reason, find all nodes without parents _and_ simply hardcode
+    # `<built-in method builtins.exec>`.
+    roots = set()
+    for item in stats.stats.items():
+        key, value = item
+        if value[4] == {}:
+            roots.add(key)
+
+    default_roots = [
+        ("~", 0, "<built-in method builtins.exec>"),
+        ("~", 0, "<built-in method exec>"),
+    ]
+    for default_root in default_roots:
+        if default_root in stats.stats:
+            roots.add(default_root)
+    roots = list(roots)
+
+    # Collect children
+    children = {key: [] for key in stats.stats.keys()}
+    for key, value in stats.stats.items():
+        _, _, _, _, parents = value
+        for parent in parents:
+            children[parent].append(key)
+
+    def populate(key, parent):
+        if parent is None:
+            _, _, selftime, cumtime, _ = stats.stats[key]
+            parent_times = {}
+        else:
+            _, _, x, _, parent_times = stats.stats[key]
+            _, _, selftime, cumtime = parent_times[parent]
+
+        # Convert the tuple key into a string
+        name = "{}::{}::{}".format(*key)
+        if len(parent_times) <= 1:
+            # Handle children
+            # merge dictionaries
+            c = [populate(child, key) for child in children[key]]
+            c.append(
+                {
+                    "text": [name + "::self", f"{selftime:.3} s"],
+                    "color": 0,
+                    "value": selftime,
+                }
+            )
+            out = {"text": [name], "color": 0, "children": c}
+        else:
+            # More than one parent; we cannot further determine the call times.
+            # Terminate the tree here.
+            if children[key]:
+                c = [
+                    {
+                        "text": [
+                            "Possible calls of",
+                            ", ".join(
+                                "{}::{}::{}".format(*child) for child in children[key]
+                            ),
+                        ],
+                        "color": 3,
+                        "value": cumtime,
+                    }
+                ]
+                out = {
+                    "text": [name],
+                    "color": 0,
+                    "children": c,
+                }
+            else:
+                out = {
+                    "text": [name, f"{cumtime:.3f}"],
+                    "color": 0,
+                    "value": cumtime,
+                }
+        return out
+
+    data = {
+        "text": ["root"],
+        "color": 0,
+        "children": [populate(root, None) for root in roots],
+    }
+    return data