"""
This module implements code highlighting of numba function annotations.
"""
from warnings import warn
warn("The pretty_annotate functionality is experimental and might change API",
FutureWarning)
def hllines(code, style):
try:
from pygments import highlight
from pygments.lexers import PythonLexer
from pygments.formatters import HtmlFormatter
except ImportError:
raise ImportError("please install the 'pygments' package")
pylex = PythonLexer()
"Given a code string, return a list of html-highlighted lines"
hf = HtmlFormatter(noclasses=True, style=style, nowrap=True)
res = highlight(code, pylex, hf)
return res.splitlines()
def htlines(code, style):
try:
from pygments import highlight
from pygments.lexers import PythonLexer
# TerminalFormatter does not support themes, Terminal256 should,
# but seem to not work.
from pygments.formatters import TerminalFormatter
except ImportError:
raise ImportError("please install the 'pygments' package")
pylex = PythonLexer()
"Given a code string, return a list of ANSI-highlighted lines"
hf = TerminalFormatter(style=style)
res = highlight(code, pylex, hf)
return res.splitlines()
def get_ansi_template():
try:
from jinja2 import Template
except ImportError:
raise ImportError("please install the 'jinja2' package")
return Template("""
{%- for func_key in func_data.keys() -%}
Function name: \x1b[34m{{func_data[func_key]['funcname']}}\x1b[39;49;00m
{%- if func_data[func_key]['filename'] -%}
{{'\n'}}In file: \x1b[34m{{func_data[func_key]['filename'] -}}\x1b[39;49;00m
{%- endif -%}
{{'\n'}}With signature: \x1b[34m{{func_key[1]}}\x1b[39;49;00m
{{- "\n" -}}
{%- for num, line, hl, hc in func_data[func_key]['pygments_lines'] -%}
{{-'\n'}}{{ num}}: {{hc-}}
{%- if func_data[func_key]['ir_lines'][num] -%}
{%- for ir_line, ir_line_type in func_data[func_key]['ir_lines'][num] %}
{{-'\n'}}--{{- ' '*func_data[func_key]['python_indent'][num]}}
{{- ' '*(func_data[func_key]['ir_indent'][num][loop.index0]+4)
}}{{ir_line }}\x1b[41m{{ir_line_type-}}\x1b[39;49;00m
{%- endfor -%}
{%- endif -%}
{%- endfor -%}
{%- endfor -%}
""")
return ansi_template
def get_html_template():
try:
from jinja2 import Template
except ImportError:
raise ImportError("please install the 'jinja2' package")
return Template("""
{% for func_key in func_data.keys() %}
Function name: {{func_data[func_key]['funcname']}}
{% if func_data[func_key]['filename'] %}
in file: {{func_data[func_key]['filename']|escape}}
{% endif %}
with signature: {{func_key[1]|e}}
{%- for num, line, hl, hc in func_data[func_key]['pygments_lines'] -%}
{%- if func_data[func_key]['ir_lines'][num] %}
{{num}}:
{{' '*func_data[func_key]['python_indent'][num]}}{{hl}}
{%- for ir_line, ir_line_type in func_data[func_key]['ir_lines'][num] %}
{{- ' '*func_data[func_key]['python_indent'][num]}}
{{ ' '*func_data[func_key]['ir_indent'][num][loop.index0]}}{{ir_line|e -}}
{{ir_line_type}}
|
{%- endfor -%}
|
{% else -%}
{{num}}:
{{' '*func_data[func_key]['python_indent'][num]}}{{hl}}
|
{%- endif -%}
{%- endfor -%}
{% endfor %}
""")
def reform_code(annotation):
"""
Extract the code from the Numba annotation datastructure.
Pygments can only highlight full multi-line strings, the Numba
annotation is list of single lines, with indentation removed.
"""
ident_dict = annotation['python_indent']
s= ''
for n,l in annotation['python_lines']:
s = s+' '*ident_dict[n]+l+'\n'
return s
class Annotate:
"""
Construct syntax highlighted annotation for a given jitted function:
Example:
>>> import numba
>>> from numba.pretty_annotate import Annotate
>>> @numba.jit
... def test(q):
... res = 0
... for i in range(q):
... res += i
... return res
...
>>> test(10)
45
>>> Annotate(test)
The last line will return an HTML and/or ANSI representation that will be
displayed accordingly in Jupyter/IPython.
Function annotations persist across compilation for newly encountered
type signatures and as a result annotations are shown for all signatures
by default.
Annotations for a specific signature can be shown by using the
``signature`` parameter.
>>> @numba.jit
... def add(x, y):
... return x + y
...
>>> add(1, 2)
3
>>> add(1.3, 5.7)
7.0
>>> add.signatures
[(int64, int64), (float64, float64)]
>>> Annotate(add, signature=add.signatures[1]) # annotation for (float64, float64)
"""
def __init__(self, function, signature=None, **kwargs):
style = kwargs.get('style', 'default')
if not function.signatures:
raise ValueError('function need to be jitted for at least one signature')
ann = function.get_annotation_info(signature=signature)
self.ann = ann
for k,v in ann.items():
res = hllines(reform_code(v), style)
rest = htlines(reform_code(v), style)
v['pygments_lines'] = [(a,b,c, d) for (a,b),c, d in zip(v['python_lines'], res, rest)]
def _repr_html_(self):
return get_html_template().render(func_data=self.ann)
def __repr__(self):
return get_ansi_template().render(func_data=self.ann)