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 94a7ba37 authored by Pierre Augier's avatar Pierre Augier
Browse files

AnnAssign for capturex

parent 68ce1b6da800
import gast as ast
from transonic.analyses.capturex import CaptureX
from transonic.analyses import extast
from transonic.analyses.util import print_dumped, print_unparsed
with open("simple.py") as file:
code = file.read()
module = extast.parse(code)
# print_dumped(module)
for node in module.body:
if isinstance(node, ast.FunctionDef):
fdef = node
break
capturex = CaptureX([fdef], module, consider_annotations="only")
code_ext = capturex.make_code_external()
print("code:", code_ext, sep="\n")
\ No newline at end of file
tmp = float
T = tmp
T1 = int
const = 1
const_ann = 1
def func():
i: T
i = const
ii: T1 = const_ann
return i * ii
......@@ -98,6 +98,15 @@ class CaptureX(ast.NodeVisitor):
self.visit(node)
self._annot = None
def visit_AnnAssign(self, node):
if self.consider_annotations:
self._annot = node.annotation
self.visit(node.annotation)
self._annot = None
if node.value is not None and self.consider_annotations != "only":
self.visit(node.value)
def rec(self, node):
"walk definitions to find their operands's def"
if node is self.func:
......
......@@ -11,7 +11,6 @@ import gast as ast
import astunparse
from transonic.analyses import extast
from transonic.analyses.capturex import CaptureX
def print_dumped(source):
......@@ -32,6 +31,9 @@ def print_unparsed(node):
print(extast.unparse(node))
from transonic.analyses.capturex import CaptureX
def _fill_ast_annotations_function(function_def, ast_annotations):
dict_node = ast_annotations.value
......
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