Read about our upcoming Code of Conduct on this issue

This instance will be upgraded to Heptapod 0.28.1 on 2022-01-26 at 16:00 UTC+1 (a few minutes of down time)

Commit e1e37aeb authored by Pierre Augier's avatar Pierre Augier
Browse files

extast: extended AST with CommentLine nodes

parent c39b40735f46
import gast as ast
import beniget
import astunparse
class CaptureX(ast.NodeVisitor):
......@@ -49,6 +48,8 @@ class CaptureX(ast.NodeVisitor):
if __name__ == "__main__":
from transonic.analyses import extast
code = "a = 1; b = [a, a]\ndef foo():\n return b"
# code = "a = 1; b = len([a, a])\ndef foo():\n return b"
# code = "import numpy as np\na = np.int(1)\ndef foo():\n return np.zeros(a)"
......@@ -58,9 +59,13 @@ if __name__ == "__main__":
a = 1
def fooo():
# comment fooo
return 1
# comment in module
def foo():
# comment in foo
return a + fooo()
def bar():
......@@ -69,11 +74,11 @@ def bar():
"""
module = ast.parse(code)
module = extast.parse(code)
function = module.body[3]
capturex = CaptureX(module, function)
capturex.visit(function)
for node in capturex.external:
# print(astunparse.dump(node))
print(astunparse.unparse(node).strip())
print(extast.unparse(node).strip())
......@@ -7,7 +7,7 @@ from .util import filter_code_typevars, get_annotations
from .capturex import CaptureX, make_code_external
from .blocks_if import get_block_definitions
from .parser import parse_code
from . import extast
from ..log import logger
......@@ -77,8 +77,8 @@ def analyse_aot(code):
debug = logger.debug
debug("ast.parse")
module = ast.parse(code)
debug("extast.parse")
module = extast.parse(code)
debug("compute ancestors and chains")
ancestors, duc, udc = compute_ancestors_chains(module)
......
import gast as ast
import beniget
import astunparse
from transonic.analyses import extast
class CaptureX(ast.NodeVisitor):
......@@ -109,7 +110,7 @@ class CaptureX(ast.NodeVisitor):
def make_code_external(external, code_module=None):
code = []
for node in external:
code.append(astunparse.unparse(node).strip())
code.append(extast.unparse(node).strip())
return "\n".join(code)
......@@ -134,7 +135,7 @@ def bar():
"""
module = ast.parse(code)
module = extast.parse(code)
function = module.body[3]
capturex = CaptureX((function,), module)
......
"""Extended AST with CommentLine nodes
======================================
"""
from io import StringIO
import gast as ast
import astunparse
import beniget
class CommentLine(ast.AST):
_fields = ("s",)
def __init__(self, s, lineno=None):
super().__init__()
self.s = s
if lineno is not None:
self.lineno = lineno
class UnparserExtended(astunparse.Unparser):
def __init__(self, tree, file, with_comments=True):
self.with_comments = with_comments
super().__init__(tree, file=file)
def _CommentLine(self, node):
if self.with_comments:
self.write(f"\n{' '* self._indent}{node.s}")
def parse(code, *args, **kwargs):
tree = ast.parse(code, *args, **kwargs)
CommentInserter(tree, code)
return tree
def unparse(tree, with_comments=True):
v = StringIO()
UnparserExtended(tree, file=v, with_comments=with_comments)
return v.getvalue()
class CommentInserter(ast.NodeVisitor):
def __init__(self, tree, code):
self.tree = tree
self.ancestors = beniget.Ancestors()
self.ancestors.visit(tree)
self.code = code
self.lines = code.split("\n")
self.lines_comments = []
self.lineno_comments = []
for index, line in enumerate(self.lines):
if line.strip().startswith("#"):
self.lineno_comments.append(index + 1)
self.lines_comments.append(line)
if not self.lineno_comments:
# nothing to do
return
self.last_line_comment = max(self.lineno_comments)
self.index_comment = 0
self.node_after = {}
self.node_before = []
self.previous_node = None
self._done = False
self.visit(tree)
self.node_before = {
self.lineno_comments[index]: node_before
for index, node_before in enumerate(self.node_before)
}
self.modify_tree()
def visit(self, node):
if self._done:
return
if hasattr(node, "lineno"):
linenos = self.lineno_comments
while node.lineno > linenos[self.index_comment]:
self.node_after[linenos[self.index_comment]] = node
if linenos[self.index_comment] == self.last_line_comment:
self._done = True
return
self.index_comment += 1
while len(self.node_after) > len(self.node_before):
self.node_before.append(self.previous_node)
self.previous_node = node
super().visit(node)
def modify_tree(self):
# todo: debug this buggy code!
for lineno, line in zip(self.lineno_comments, self.lines_comments):
try:
node_after = self.node_after[lineno]
except KeyError:
self.tree.body.append(CommentLine(line, lineno))
continue
else:
self.insert_comment_in_parent_before_node(
line, lineno, node_after
)
continue
try:
node_before = self.node_before[lineno]
except KeyError:
pass
else:
if isinstance(node_before, ast.Module):
self.tree.body.insert(0, CommentLine(line, lineno))
def insert_comment_in_parent_before_node(self, line, lineno, node):
parent = self.ancestors.parent(node)
comment = line.strip()
for field, value in ast.iter_fields(parent):
if isinstance(value, list):
if node in value:
index = value.index(node)
value.insert(index, CommentLine(comment, lineno))
if __name__ == "__main__":
code = """
# comment 0
# comment 1
if True:
# comment 1bis
# comment 2
# comment 3
var_if = 1
# comment after var_if
if False:
var_ifif = 1
# comment ifif
# comment if end
# comment 4
else:
# comment else
var_else = 1
# comment after else
a += 1
# comment after last
"""
code = """
from transonic import Transonic
import numpy as np
import foo
ts = Transonic()
a = n = 1
b = np.ones(2)
if ts.is_transpiled:
result = ts.use_block("block0")
else:
# transonic block (
# A a; A1 b;
# float n
# )
# transonic block (
# int[:] a, b;
# float n
# )
result = a ** 2 + b.mean() ** 3 + foo.bar(n)
"""
tree = parse(code)
print(unparse(tree, with_comments=True))
from transonic.analyses import compute_ancestors_chains
compute_ancestors_chains(tree)
from textwrap import dedent
import gast as ast
import astunparse
from transonic.analyses import extast
def print_dumped(source):
if isinstance(source, str):
module = ast.parse(source)
module = extast.parse(source)
if len(module.body) == 1:
node = module.body[0]
else:
node = module
else:
node = source
print(astunparse.dump(node))
print(ast.dump(node))
def print_unparsed(node):
print(astunparse.unparse(node))
print(extast.unparse(node))
def _fill_ast_annotations_function(function_def, ast_annotations):
......@@ -65,7 +65,7 @@ def get_annotations(object_def, namespace):
# print_dump(ast_annotations)
source = astunparse.unparse(ast_annotations)
source = extast.unparse(ast_annotations)
try:
del namespace["__builtins__"]
......@@ -105,7 +105,7 @@ def filter_code_typevars(module, duc, ancestors):
elif isinstance(node, (ast.Assign, ast.AugAssign)):
kept.append(node)
return astunparse.unparse(module_filtered)
return extast.unparse(module_filtered)
class AnalyseLines(ast.NodeVisitor):
......
from tokenize import tokenize, untokenize, NAME, OP
from io import BytesIO
from token import tok_name
import re
from textwrap import indent
import astunparse
try:
import black
......@@ -16,29 +13,22 @@ import transonic
from transonic.compat import open
from transonic.log import logger
from transonic.util import TypeHintRemover
from transonic.analyses.util import gather_rawcode_comments
from transonic.annotation import compute_pythran_types_from_valued_types
from transonic.analyses import analyse_aot
from transonic.analyses import extast
def get_code_function(fdef, code_module):
transformed = TypeHintRemover().visit(fdef)
# convert the AST back to source code
striped_code = astunparse.unparse(transformed)
raw_code, comments = gather_rawcode_comments(fdef, code_module)
code = extast.unparse(transformed)
if black:
raw_code = black.format_str(raw_code, line_length=82)
# bad hack to conserve the comments...
sep = ":(\n)+ "
head = re.split(sep, striped_code, maxsplit=1)[0]
body = re.split(sep, raw_code, maxsplit=1)[-1]
code = black.format_str(code, line_length=82)
return head + ":\n " + body
return code
def produce_new_code_method(fdef, class_def, code_module):
......@@ -180,7 +170,7 @@ def produce_code_for_method(
if ind < nb_no_defaults:
str_args_value_func.append(f"{name}")
else:
default = astunparse.unparse(defaults[ind_default]).strip()
default = extast.unparse(defaults[ind_default]).strip()
str_args_value_func.append(f"{name}={default}")
ind_default += 1
......@@ -288,7 +278,7 @@ def make_pythran_code(path_py):
code.append(f"\ndef {block.name}({str_variables}):\n")
code_block = indent(astunparse.unparse(block.ast_code), " ")
code_block = indent(extast.unparse(block.ast_code), " ")
code.append(code_block)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment