[mypyc] Add prototype feature to generate annotated html of compiled code (#18828)

Now you can use `mypyc -a x.html foo.py` to compile `foo.py` and
generate `x.html`, which contains annotated source code for `foo.py`
with some inefficient operations highlighted with a red background.

Right now this is close to minimal and only detects a few inefficient
operations. I will improve this in follow-up PRs.

The overall idea is similar to the Cython `-a` flag or `annotate=True`.
Here is an example (scroll down a bit):

https://cython.readthedocs.io/en/latest/src/tutorial/cython_tutorial.html#primes

The approach here differs in some ways from Cython. First, we only
generate a single html file with annotations for all compiled files. I
think this will make things easier when compiling a large number of
modules. We'll probably need to add some navigation aids to the
generated html eventually. Second, instead of showing the C code when
encountering inefficient operations, we will generate (hopefully)
easy-to-understand hints that explain what is slow (and perhaps why),
without requiring any understanding of C.

Author: JukkaL

mypy/main.py

@@ -1123,6 +1123,11 @@ def add_invertible_flag(
                 dest=f"special-opts:{report_type}_report",
             )
 
+    # Undocumented mypyc feature: generate annotated HTML source file
+    report_group.add_argument(
+        "-a", dest="mypyc_annotation_file", type=str, default=None, help=argparse.SUPPRESS
+    )
+
     other_group = parser.add_argument_group(title="Miscellaneous")
     other_group.add_argument("--quickstart-file", help=argparse.SUPPRESS)
     other_group.add_argument("--junit-xml", help="Write junit.xml to the given file")

mypy/options.py

@@ -406,6 +406,9 @@ def __init__(self) -> None:
         # Sets custom output format
         self.output: str | None = None
 
+        # Output html file for mypyc -a
+        self.mypyc_annotation_file: str | None = None
+
     def use_lowercase_names(self) -> bool:
         if self.python_version >= (3, 9):
             return not self.force_uppercase_builtins

mypyc/annotate.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import os.path
+import sys
+from html import escape
+
+from mypy.build import BuildResult
+from mypy.nodes import MypyFile
+from mypy.util import FancyFormatter
+from mypyc.ir.func_ir import FuncIR
+from mypyc.ir.module_ir import ModuleIR
+from mypyc.ir.ops import CallC, LoadLiteral, Value
+
+CSS = """\
+.collapsible {
+    cursor: pointer;
+}
+
+.content {
+    display: block;
+    margin-top: 10px;
+    margin-bottom: 10px;
+}
+
+.hint {
+    display: inline;
+    border: 1px solid #ccc;
+    padding: 5px;
+}
+"""
+
+JS = """\
+document.querySelectorAll('.collapsible').forEach(function(collapsible) {
+    collapsible.addEventListener('click', function() {
+        const content = this.nextElementSibling;
+        if (content.style.display === 'none') {
+            content.style.display = 'block';
+        } else {
+            content.style.display = 'none';
+        }
+    });
+});
+"""
+
+
+class AnnotatedSource:
+    def __init__(self, path: str, annotations: dict[int, list[str]]) -> None:
+        self.path = path
+        self.annotations = annotations
+
+
+def generate_annotated_html(
+    html_fnam: str, result: BuildResult, modules: dict[str, ModuleIR]
+) -> None:
+    annotations = []
+    for mod, mod_ir in modules.items():
+        path = result.graph[mod].path
+        tree = result.graph[mod].tree
+        assert tree is not None
+        annotations.append(generate_annotations(path or "<source>", tree, mod_ir))
+    html = generate_html_report(annotations)
+    with open(html_fnam, "w") as f:
+        f.write(html)
+
+    formatter = FancyFormatter(sys.stdout, sys.stderr, False)
+    formatted = formatter.style(os.path.abspath(html_fnam), "none", underline=True, bold=True)
+    print(f"\nWrote {formatted} -- open in browser to view\n")
+
+
+def generate_annotations(path: str, tree: MypyFile, ir: ModuleIR) -> AnnotatedSource:
+    anns = {}
+    for func_ir in ir.functions:
+        anns.update(function_annotations(func_ir))
+    return AnnotatedSource(path, anns)
+
+
+def function_annotations(func_ir: FuncIR) -> dict[int, list[str]]:
+    # TODO: check if func_ir.line is -1
+    anns: dict[int, list[str]] = {}
+    for block in func_ir.blocks:
+        for op in block.ops:
+            if isinstance(op, CallC):
+                name = op.function_name
+                ann = None
+                if name == "CPyObject_GetAttr":
+                    attr_name = get_str_literal(op.args[1])
+                    if attr_name:
+                        ann = f'Get non-native attribute "{attr_name}".'
+                    else:
+                        ann = "Dynamic attribute lookup."
+                elif name == "PyNumber_Add":
+                    ann = 'Generic "+" operation.'
+                if ann:
+                    anns.setdefault(op.line, []).append(ann)
+    return anns
+
+
+def get_str_literal(v: Value) -> str | None:
+    if isinstance(v, LoadLiteral) and isinstance(v.value, str):
+        return v.value
+    return None
+
+
+def generate_html_report(sources: list[AnnotatedSource]) -> str:
+    html = []
+    html.append("<html>\n<head>\n")
+    html.append(f"<style>\n{CSS}\n</style>")
+    html.append("</head>\n")
+    html.append("<body>\n")
+    for src in sources:
+        html.append(f"<h2><tt>{src.path}</tt></h2>\n")
+        html.append("<pre>")
+        anns = src.annotations
+        with open(src.path) as f:
+            lines = f.readlines()
+        for i, s in enumerate(lines):
+            s = escape(s)
+            line = i + 1
+            linenum = "%5d" % line
+            if line in anns:
+                hint = " ".join(anns[line])
+                s = colorize_line(linenum, s, hint_html=hint)
+            else:
+                s = linenum + "  " + s
+            html.append(s)
+        html.append("</pre>")
+
+    html.append("<script>")
+    html.append(JS)
+    html.append("</script>")
+
+    html.append("</body></html>\n")
+    return "".join(html)
+
+
+def colorize_line(linenum: str, s: str, hint_html: str) -> str:
+    hint_prefix = " " * len(linenum) + "  "
+    line_span = f'<div class="collapsible" style="background-color: #fcc">{linenum}  {s}</div>'
+    hint_div = f'<div class="content">{hint_prefix}<div class="hint">{hint_html}</div></div>'
+    return f"<span>{line_span}{hint_div}</span>"

mypyc/build.py

@@ -34,6 +34,7 @@
 from mypy.main import process_options
 from mypy.options import Options
 from mypy.util import write_junit_xml
+from mypyc.annotate import generate_annotated_html
 from mypyc.codegen import emitmodule
 from mypyc.common import RUNTIME_C_FILES, shared_lib_name
 from mypyc.errors import Errors
@@ -253,6 +254,9 @@ def generate_c(
     if compiler_options.verbose:
         print(f"Compiled to C in {t2 - t1:.3f}s")
 
+    if options.mypyc_annotation_file:
+        generate_annotated_html(options.mypyc_annotation_file, result, modules)
+
     return ctext, "\n".join(format_modules(modules))
 
 

mypyc/test-data/annotate-basic.test

@@ -0,0 +1,26 @@
+[case testAnnotateNonNativeAttribute]
+def f(x):
+    return x.foo
+
+class C:
+    foo: int
+
+def g(x: C) -> int:
+    return x.foo
+[out]
+2: Get non-native attribute "foo".
+
+[case testAnnotateGenericAdd]
+def f(x):
+    return x + 1
+
+def g(x: int) -> int:
+    return x + 1
+[out]
+2: Generic "+" operation.
+
+[case testAnnotateTwoOperationsOnLine]
+def f(x):
+    return x.foo + 1
+[out]
+2: Get non-native attribute "foo". Generic "+" operation.

mypyc/test/test_alwaysdefined.py

@@ -31,7 +31,7 @@ def run_case(self, testcase: DataDrivenTestCase) -> None:
             return
         with use_custom_builtins(os.path.join(self.data_prefix, ICODE_GEN_BUILTINS), testcase):
             try:
-                ir = build_ir_for_single_file2(testcase.input, options)
+                ir = build_ir_for_single_file2(testcase.input, options)[0]
             except CompileError as e:
                 actual = e.messages
             else:

mypyc/test/test_annotate.py

@@ -0,0 +1,48 @@
+"""Test cases for annotating source code to highlight inefficiencies."""
+
+from __future__ import annotations
+
+import os.path
+
+from mypy.errors import CompileError
+from mypy.test.config import test_temp_dir
+from mypy.test.data import DataDrivenTestCase
+from mypyc.annotate import generate_annotations
+from mypyc.test.testutil import (
+    ICODE_GEN_BUILTINS,
+    MypycDataSuite,
+    assert_test_output,
+    build_ir_for_single_file2,
+    infer_ir_build_options_from_test_name,
+    remove_comment_lines,
+    use_custom_builtins,
+)
+
+files = ["annotate-basic.test"]
+
+
+class TestReport(MypycDataSuite):
+    files = files
+    base_path = test_temp_dir
+    optional_out = True
+
+    def run_case(self, testcase: DataDrivenTestCase) -> None:
+        """Perform a runtime checking transformation test case."""
+        options = infer_ir_build_options_from_test_name(testcase.name)
+        if options is None:
+            # Skipped test case
+            return
+        with use_custom_builtins(os.path.join(self.data_prefix, ICODE_GEN_BUILTINS), testcase):
+            expected_output = remove_comment_lines(testcase.output)
+            try:
+                ir, tree = build_ir_for_single_file2(testcase.input, options)
+            except CompileError as e:
+                actual = e.messages
+            else:
+                annotations = generate_annotations("native.py", tree, ir)
+                actual = []
+                for line, line_anns in annotations.annotations.items():
+                    s = " ".join(line_anns)
+                    actual.append(f"{line}: {s}")
+
+            assert_test_output(testcase, actual, "Invalid source code output", expected_output)

mypyc/test/testutil.py

@@ -12,6 +12,7 @@
 
 from mypy import build
 from mypy.errors import CompileError
+from mypy.nodes import MypyFile
 from mypy.options import Options
 from mypy.test.config import test_temp_dir
 from mypy.test.data import DataDrivenTestCase, DataSuite
@@ -93,12 +94,12 @@ def perform_test(
 def build_ir_for_single_file(
     input_lines: list[str], compiler_options: CompilerOptions | None = None
 ) -> list[FuncIR]:
-    return build_ir_for_single_file2(input_lines, compiler_options).functions
+    return build_ir_for_single_file2(input_lines, compiler_options)[0].functions
 
 
 def build_ir_for_single_file2(
     input_lines: list[str], compiler_options: CompilerOptions | None = None
-) -> ModuleIR:
+) -> tuple[ModuleIR, MypyFile]:
     program_text = "\n".join(input_lines)
 
     # By default generate IR compatible with the earliest supported Python C API.
@@ -137,7 +138,9 @@ def build_ir_for_single_file2(
     module = list(modules.values())[0]
     for fn in module.functions:
         assert_func_ir_valid(fn)
-    return module
+    tree = result.graph[module.fullname].tree
+    assert tree is not None
+    return module, tree
 
 
 def update_testcase_output(testcase: DataDrivenTestCase, output: list[str]) -> None: