"""Generate source code formatted as HTML, with bottlenecks annotated and highlighted.

Various heuristics are used to detect common issues that cause slower than

expected performance.

"""

from __future__ import annotations

import os.path

import sys

from html import escape

from typing import Final

from mypy.build import BuildResult

from mypy.nodes import (

AssignmentStmt,

CallExpr,

ClassDef,

Decorator,

DictionaryComprehension,

Expression,

ForStmt,

FuncDef,

GeneratorExpr,

IndexExpr,

LambdaExpr,

MemberExpr,

MypyFile,

NamedTupleExpr,

NameExpr,

NewTypeExpr,

Node,

OpExpr,

RefExpr,

TupleExpr,

TypedDictExpr,

TypeInfo,

TypeVarExpr,

Var,

WithStmt,

)

from mypy.traverser import TraverserVisitor

from mypy.types import AnyType, Instance, ProperType, Type, TypeOfAny, get_proper_type

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, LoadStatic, Value

from mypyc.irbuild.mapper import Mapper

class Annotation:

"""HTML annotation for compiled source code"""

def __init__(self, message: str, priority: int = 1) -> None:

# Message as HTML that describes an issue and/or how to fix it.

# Multiple messages on a line may be concatenated.

self.message = message

# If multiple annotations are generated for a single line, only report

# the highest-priority ones. Some use cases generate multiple annotations,

# and this can be used to reduce verbosity by hiding the lower-priority

# ones.

self.priority = priority

op_hints: Final = {

"PyNumber_Add": Annotation('Generic "+" operation.'),

"PyNumber_Subtract": Annotation('Generic "-" operation.'),

"PyNumber_Multiply": Annotation('Generic "*" operation.'),

"PyNumber_TrueDivide": Annotation('Generic "/" operation.'),

"PyNumber_FloorDivide": Annotation('Generic "//" operation.'),

"PyNumber_Positive": Annotation('Generic unary "+" operation.'),

"PyNumber_Negative": Annotation('Generic unary "-" operation.'),

"PyNumber_And": Annotation('Generic "&" operation.'),

"PyNumber_Or": Annotation('Generic "|" operation.'),

"PyNumber_Xor": Annotation('Generic "^" operation.'),

"PyNumber_Lshift": Annotation('Generic "<<" operation.'),

"PyNumber_Rshift": Annotation('Generic ">>" operation.'),

"PyNumber_Invert": Annotation('Generic "~" operation.'),

"PyObject_Call": Annotation("Generic call operation."),

"PyObject_CallObject": Annotation("Generic call operation."),

"PyObject_RichCompare": Annotation("Generic comparison operation."),

"PyObject_GetItem": Annotation("Generic indexing operation."),

"PyObject_SetItem": Annotation("Generic indexed assignment."),

}

stdlib_hints: Final = {

"functools.partial": Annotation(

'"functools.partial" is inefficient in compiled code.', priority=3

),

"itertools.chain": Annotation(

'"itertools.chain" is inefficient in compiled code (hint: replace with for loops).',

priority=3,

),

"itertools.groupby": Annotation(

'"itertools.groupby" is inefficient in compiled code.', priority=3

),

"itertools.islice": Annotation(

'"itertools.islice" is inefficient in compiled code (hint: replace with for loop over index range).',

priority=3,

),

"copy.deepcopy": Annotation(

'"copy.deepcopy" tends to be slow. Make a shallow copy if possible.', priority=2

),

}

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:

"""Annotations for a single compiled source file."""

def __init__(self, path: str, annotations: dict[int, list[Annotation]]) -> None:

self.path = path

self.annotations = annotations

def generate_annotated_html(

html_fnam: str, result: BuildResult, modules: dict[str, ModuleIR], mapper: Mapper

) -> 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, result.types, mapper)

)

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, type_map: dict[Expression, Type], mapper: Mapper

) -> AnnotatedSource:

anns = {}

for func_ir in ir.functions:

anns.update(function_annotations(func_ir, tree))

visitor = ASTAnnotateVisitor(type_map, mapper)

for defn in tree.defs:

defn.accept(visitor)

anns.update(visitor.anns)

for line in visitor.ignored_lines:

if line in anns:

del anns[line]

return AnnotatedSource(path, anns)

def function_annotations(func_ir: FuncIR, tree: MypyFile) -> dict[int, list[Annotation]]:

"""Generate annotations based on mypyc IR."""

# TODO: check if func_ir.line is -1

anns: dict[int, list[Annotation]] = {}

for block in func_ir.blocks:

for op in block.ops:

if isinstance(op, CallC):

name = op.function_name

ann: str | Annotation | None = None

if name == "CPyObject_GetAttr":

attr_name = get_str_literal(op.args[1])

if attr_name in ("__prepare__", "GeneratorExit", "StopIteration"):

# These attributes are internal to mypyc/CPython, and/or accessed

# implicitly in generated code. The user has little control over

# them.

ann = None

elif attr_name:

ann = f'Get non-native attribute "{attr_name}".'

else:

ann = "Dynamic attribute lookup."

elif name == "PyObject_SetAttr":

attr_name = get_str_literal(op.args[1])

if attr_name == "__mypyc_attrs__":

# This is set implicitly and can't be avoided.

ann = None

elif attr_name:

ann = f'Set non-native attribute "{attr_name}".'

else:

ann = "Dynamic attribute set."

elif name == "PyObject_VectorcallMethod":

method_name = get_str_literal(op.args[0])

if method_name:

ann = f'Call non-native method "{method_name}" (it may be defined in a non-native class, or decorated).'

else:

ann = "Dynamic method call."

elif name in op_hints:

ann = op_hints[name]

elif name in ("CPyDict_GetItem", "CPyDict_SetItem"):

if (

isinstance(op.args[0], LoadStatic)

and isinstance(op.args[1], LoadLiteral)

and func_ir.name != "__top_level__"

):

load = op.args[0]

name = str(op.args[1].value)

sym = tree.names.get(name)

if (

sym

and sym.node

and load.namespace == "static"

and load.identifier == "globals"

):

if sym.node.fullname in stdlib_hints:

ann = stdlib_hints[sym.node.fullname]

elif isinstance(sym.node, Var):

ann = (

f'Access global "{name}" through namespace '

+ "dictionary (hint: access is faster if you can make it Final)."

)

else:

ann = f'Access "{name}" through global namespace dictionary.'

if ann:

if isinstance(ann, str):

ann = Annotation(ann)

anns.setdefault(op.line, []).append(ann)

return anns

class ASTAnnotateVisitor(TraverserVisitor):

"""Generate annotations from mypy AST and inferred types."""

def __init__(self, type_map: dict[Expression, Type], mapper: Mapper) -> None:

self.anns: dict[int, list[Annotation]] = {}

self.ignored_lines: set[int] = set()

self.func_depth = 0

self.type_map = type_map

self.mapper = mapper

def visit_func_def(self, o: FuncDef, /) -> None:

if self.func_depth > 0:

self.annotate(

o,

"A nested function object is allocated each time statement is executed. "

+ "A module-level function would be faster.",

)

self.func_depth += 1

super().visit_func_def(o)

self.func_depth -= 1

def visit_for_stmt(self, o: ForStmt, /) -> None:

self.check_iteration([o.expr], "For loop")

super().visit_for_stmt(o)

def visit_dictionary_comprehension(self, o: DictionaryComprehension, /) -> None:

self.check_iteration(o.sequences, "Comprehension")

super().visit_dictionary_comprehension(o)

def visit_generator_expr(self, o: GeneratorExpr, /) -> None:

self.check_iteration(o.sequences, "Comprehension or generator")

super().visit_generator_expr(o)

def check_iteration(self, expressions: list[Expression], kind: str) -> None:

for expr in expressions:

typ = self.get_type(expr)

if isinstance(typ, AnyType):

self.annotate(expr, f'{kind} uses generic operations (iterable has type "Any").')

elif isinstance(typ, Instance) and typ.type.fullname in (

"typing.Iterable",

"typing.Iterator",

"typing.Sequence",

"typing.MutableSequence",

):

self.annotate(

expr,

f'{kind} uses generic operations (iterable has the abstract type "{typ.type.fullname}").',

)

def visit_class_def(self, o: ClassDef, /) -> None:

super().visit_class_def(o)

if self.func_depth == 0:

# Don't complain about base classes at top level

for base in o.base_type_exprs:

self.ignored_lines.add(base.line)

for s in o.defs.body:

if isinstance(s, AssignmentStmt):

# Don't complain about attribute initializers

self.ignored_lines.add(s.line)

elif isinstance(s, Decorator):

# Don't complain about decorator definitions that generate some

# dynamic operations. This is a bit heavy-handed.

self.ignored_lines.add(s.func.line)

def visit_with_stmt(self, o: WithStmt, /) -> None:

for expr in o.expr:

if isinstance(expr, CallExpr) and isinstance(expr.callee, RefExpr):

node = expr.callee.node

if isinstance(node, Decorator):

if any(

isinstance(d, RefExpr)

and d.node

and d.node.fullname == "contextlib.contextmanager"

for d in node.decorators

):

self.annotate(

expr,

f'"{node.name}" uses @contextmanager, which is slow '

+ "in compiled code. Use a native class with "

+ '"__enter__" and "__exit__" methods instead.',

priority=3,

)

super().visit_with_stmt(o)

def visit_assignment_stmt(self, o: AssignmentStmt, /) -> None:

special_form = False

if self.func_depth == 0:

analyzed: Expression | None = o.rvalue

if isinstance(o.rvalue, (CallExpr, IndexExpr, OpExpr)):

analyzed = o.rvalue.analyzed

if o.is_alias_def or isinstance(

analyzed, (TypeVarExpr, NamedTupleExpr, TypedDictExpr, NewTypeExpr)

):

special_form = True

if special_form:

# TODO: Ignore all lines if multi-line

self.ignored_lines.add(o.line)

super().visit_assignment_stmt(o)

def visit_name_expr(self, o: NameExpr, /) -> None:

if ann := stdlib_hints.get(o.fullname):

self.annotate(o, ann)

def visit_member_expr(self, o: MemberExpr, /) -> None:

super().visit_member_expr(o)

if ann := stdlib_hints.get(o.fullname):

self.annotate(o, ann)

def visit_call_expr(self, o: CallExpr, /) -> None:

super().visit_call_expr(o)

if (

isinstance(o.callee, RefExpr)

and o.callee.fullname == "builtins.isinstance"

and len(o.args) == 2

):

arg = o.args[1]

self.check_isinstance_arg(arg)

elif isinstance(o.callee, RefExpr) and isinstance(o.callee.node, TypeInfo):

info = o.callee.node

class_ir = self.mapper.type_to_ir.get(info)

if (class_ir and not class_ir.is_ext_class) or (

class_ir is None and not info.fullname.startswith("builtins.")

):

self.annotate(

o, f'Creating an instance of non-native class "{info.name}" ' + "is slow.", 2

)

elif class_ir and class_ir.is_augmented:

self.annotate(

o,

f'Class "{info.name}" is only partially native, and '

+ "constructing an instance is slow.",

2,

)

elif isinstance(o.callee, RefExpr) and isinstance(o.callee.node, Decorator):

decorator = o.callee.node

if self.mapper.is_native_ref_expr(o.callee):

self.annotate(

o,

f'Calling a decorated function ("{decorator.name}") is inefficient, even if it\'s native.',

2,

)

def check_isinstance_arg(self, arg: Expression) -> None:

if isinstance(arg, RefExpr):

if isinstance(arg.node, TypeInfo) and arg.node.is_protocol:

self.annotate(

arg, f'Expensive isinstance() check against protocol "{arg.node.name}".'

)

elif isinstance(arg, TupleExpr):

for item in arg.items:

self.check_isinstance_arg(item)

def visit_lambda_expr(self, o: LambdaExpr, /) -> None:

self.annotate(

o,

"A new object is allocated for lambda each time it is evaluated. "

+ "A module-level function would be faster.",

)

super().visit_lambda_expr(o)

def annotate(self, o: Node, ann: str | Annotation, priority: int = 1) -> None:

if isinstance(ann, str):

ann = Annotation(ann, priority=priority)

self.anns.setdefault(o.line, []).append(ann)

def get_type(self, e: Expression) -> ProperType:

t = self.type_map.get(e)

if t:

return get_proper_type(t)

return AnyType(TypeOfAny.unannotated)

def get_str_literal(v: Value) -> str | None:

if isinstance(v, LoadLiteral) and isinstance(v.value, str):

return v.value

return None

def get_max_prio(anns: list[Annotation]) -> list[Annotation]:

max_prio = max(a.priority for a in anns)

return [a for a in anns if a.priority == max_prio]

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>")

src_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 src_anns:

anns = get_max_prio(src_anns[line])

ann_strs = [a.message for a in anns]

hint = " ".join(ann_strs)

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>"