mirror of https://github.com/m-labs/artiq
528 lines
20 KiB
Python
528 lines
20 KiB
Python
"""
|
|
:class:`ASTTypedRewriter` rewrites a parsetree (:mod:`pythonparser.ast`)
|
|
to a typedtree (:mod:`..asttyped`).
|
|
"""
|
|
|
|
from collections import OrderedDict
|
|
from pythonparser import ast, algorithm, diagnostic
|
|
from .. import asttyped, types, builtins
|
|
|
|
# This visitor will be called for every node with a scope,
|
|
# i.e.: class, function, comprehension, lambda
|
|
class LocalExtractor(algorithm.Visitor):
|
|
def __init__(self, env_stack, engine):
|
|
super().__init__()
|
|
self.env_stack = env_stack
|
|
self.engine = engine
|
|
|
|
self.in_root = False
|
|
self.in_assign = False
|
|
self.typing_env = OrderedDict()
|
|
|
|
# which names are global have to be recorded in the current scope
|
|
self.global_ = set()
|
|
|
|
# which names are nonlocal only affects whether the current scope
|
|
# gets a new binding or not, so we throw this away
|
|
self.nonlocal_ = set()
|
|
|
|
# parameters can't be declared as global or nonlocal
|
|
self.params = set()
|
|
|
|
def visit_in_assign(self, node, in_assign):
|
|
try:
|
|
old_in_assign, self.in_assign = self.in_assign, in_assign
|
|
return self.visit(node)
|
|
finally:
|
|
self.in_assign = old_in_assign
|
|
|
|
def visit_Assign(self, node):
|
|
self.visit(node.value)
|
|
self.visit_in_assign(node.targets, in_assign=True)
|
|
|
|
def visit_For(self, node):
|
|
self.visit(node.iter)
|
|
self.visit_in_assign(node.target, in_assign=True)
|
|
self.visit(node.body)
|
|
self.visit(node.orelse)
|
|
|
|
def visit_withitem(self, node):
|
|
self.visit(node.context_expr)
|
|
self.visit_in_assign(node.optional_vars, in_assign=True)
|
|
|
|
def visit_comprehension(self, node):
|
|
self.visit(node.iter)
|
|
self.visit_in_assign(node.target, in_assign=True)
|
|
self.visit(node.ifs)
|
|
|
|
def visit_generator(self, node):
|
|
if self.in_root:
|
|
return
|
|
self.in_root = True
|
|
self.visit(list(reversed(node.generators)))
|
|
self.visit(node.elt)
|
|
|
|
visit_ListComp = visit_generator
|
|
visit_SetComp = visit_generator
|
|
visit_GeneratorExp = visit_generator
|
|
|
|
def visit_DictComp(self, node):
|
|
if self.in_root:
|
|
return
|
|
self.in_root = True
|
|
self.visit(list(reversed(node.generators)))
|
|
self.visit(node.key)
|
|
self.visit(node.value)
|
|
|
|
def visit_root(self, node):
|
|
if self.in_root:
|
|
return
|
|
self.in_root = True
|
|
self.generic_visit(node)
|
|
|
|
visit_Module = visit_root # don't look at inner scopes
|
|
visit_ClassDef = visit_root
|
|
visit_Lambda = visit_root
|
|
|
|
def visit_FunctionDef(self, node):
|
|
if self.in_root:
|
|
self._assignable(node.name)
|
|
self.visit_root(node)
|
|
|
|
def _assignable(self, name):
|
|
assert name is not None
|
|
if name not in self.typing_env and name not in self.nonlocal_:
|
|
self.typing_env[name] = types.TVar()
|
|
|
|
def visit_arg(self, node):
|
|
if node.arg in self.params:
|
|
diag = diagnostic.Diagnostic("error",
|
|
"duplicate parameter '{name}'", {"name": node.arg},
|
|
node.loc)
|
|
self.engine.process(diag)
|
|
return
|
|
self._assignable(node.arg)
|
|
self.params.add(node.arg)
|
|
|
|
def visit_Name(self, node):
|
|
if self.in_assign:
|
|
# code like:
|
|
# x = 1
|
|
# def f():
|
|
# x = 1
|
|
# creates a new binding for x in f's scope
|
|
self._assignable(node.id)
|
|
|
|
def visit_Attribute(self, node):
|
|
self.visit_in_assign(node.value, in_assign=False)
|
|
|
|
def visit_Subscript(self, node):
|
|
self.visit_in_assign(node.value, in_assign=False)
|
|
self.visit_in_assign(node.slice, in_assign=False)
|
|
|
|
def _check_not_in(self, name, names, curkind, newkind, loc):
|
|
if name in names:
|
|
diag = diagnostic.Diagnostic("error",
|
|
"name '{name}' cannot be {curkind} and {newkind} simultaneously",
|
|
{"name": name, "curkind": curkind, "newkind": newkind}, loc)
|
|
self.engine.process(diag)
|
|
return True
|
|
return False
|
|
|
|
def visit_Global(self, node):
|
|
for name, loc in zip(node.names, node.name_locs):
|
|
if self._check_not_in(name, self.nonlocal_, "nonlocal", "global", loc) or \
|
|
self._check_not_in(name, self.params, "a parameter", "global", loc):
|
|
continue
|
|
|
|
self.global_.add(name)
|
|
if len(self.env_stack) == 1:
|
|
self._assignable(name) # already in global scope
|
|
else:
|
|
if name not in self.env_stack[1]:
|
|
self.env_stack[1][name] = types.TVar()
|
|
self.typing_env[name] = self.env_stack[1][name]
|
|
|
|
def visit_Nonlocal(self, node):
|
|
for name, loc in zip(node.names, node.name_locs):
|
|
if self._check_not_in(name, self.global_, "global", "nonlocal", loc) or \
|
|
self._check_not_in(name, self.params, "a parameter", "nonlocal", loc):
|
|
continue
|
|
|
|
# nonlocal does not search prelude and global scopes
|
|
found = False
|
|
for outer_env in reversed(self.env_stack[2:]):
|
|
if name in outer_env:
|
|
found = True
|
|
break
|
|
if not found:
|
|
diag = diagnostic.Diagnostic("error",
|
|
"cannot declare name '{name}' as nonlocal: it is not bound in any outer scope",
|
|
{"name": name},
|
|
loc, [node.keyword_loc])
|
|
self.engine.process(diag)
|
|
continue
|
|
|
|
self.nonlocal_.add(name)
|
|
|
|
def visit_ExceptHandler(self, node):
|
|
self.visit(node.type)
|
|
if node.name is not None:
|
|
self._assignable(node.name)
|
|
for stmt in node.body:
|
|
self.visit(stmt)
|
|
|
|
|
|
class ASTTypedRewriter(algorithm.Transformer):
|
|
"""
|
|
:class:`ASTTypedRewriter` converts an untyped AST to a typed AST
|
|
where all type fields of non-literals are filled with fresh type variables,
|
|
and type fields of literals are filled with corresponding types.
|
|
|
|
:class:`ASTTypedRewriter` also discovers the scope of variable bindings
|
|
via :class:`LocalExtractor`.
|
|
"""
|
|
|
|
def __init__(self, engine, prelude):
|
|
self.engine = engine
|
|
self.globals = None
|
|
self.env_stack = [prelude]
|
|
self.in_class = None
|
|
|
|
def _try_find_name(self, name):
|
|
if self.in_class is not None:
|
|
typ = self.in_class.constructor_type.attributes.get(name)
|
|
if typ is not None:
|
|
return typ
|
|
|
|
for typing_env in reversed(self.env_stack):
|
|
if name in typing_env:
|
|
return typing_env[name]
|
|
|
|
def _find_name(self, name, loc):
|
|
typ = self._try_find_name(name)
|
|
if typ is not None:
|
|
return typ
|
|
|
|
diag = diagnostic.Diagnostic("fatal",
|
|
"undefined variable '{name}'", {"name":name}, loc)
|
|
self.engine.process(diag)
|
|
|
|
# Visitors that replace node with a typed node
|
|
#
|
|
def visit_Module(self, node):
|
|
extractor = LocalExtractor(env_stack=self.env_stack, engine=self.engine)
|
|
extractor.visit(node)
|
|
|
|
node = asttyped.ModuleT(
|
|
typing_env=extractor.typing_env, globals_in_scope=extractor.global_,
|
|
body=node.body, loc=node.loc)
|
|
self.globals = node.typing_env
|
|
|
|
try:
|
|
self.env_stack.append(node.typing_env)
|
|
return self.generic_visit(node)
|
|
finally:
|
|
self.env_stack.pop()
|
|
|
|
def visit_FunctionDef(self, node):
|
|
extractor = LocalExtractor(env_stack=self.env_stack, engine=self.engine)
|
|
extractor.visit(node)
|
|
|
|
signature_type = self._find_name(node.name, node.name_loc)
|
|
|
|
node = asttyped.FunctionDefT(
|
|
typing_env=extractor.typing_env, globals_in_scope=extractor.global_,
|
|
signature_type=signature_type, return_type=types.TVar(),
|
|
name=node.name, args=node.args, returns=node.returns,
|
|
body=node.body, decorator_list=node.decorator_list,
|
|
keyword_loc=node.keyword_loc, name_loc=node.name_loc,
|
|
arrow_loc=node.arrow_loc, colon_loc=node.colon_loc, at_locs=node.at_locs,
|
|
loc=node.loc, remote_fn=False)
|
|
|
|
try:
|
|
self.env_stack.append(node.typing_env)
|
|
return self.generic_visit(node)
|
|
finally:
|
|
self.env_stack.pop()
|
|
|
|
def visit_ClassDef(self, node):
|
|
if any(node.bases) or any(node.keywords) or \
|
|
node.starargs is not None or node.kwargs is not None:
|
|
diag = diagnostic.Diagnostic("error",
|
|
"inheritance is not supported", {},
|
|
node.lparen_loc.join(node.rparen_loc))
|
|
self.engine.process(diag)
|
|
|
|
for child in node.body:
|
|
if isinstance(child, (ast.Assign, ast.FunctionDef, ast.Pass)):
|
|
continue
|
|
|
|
diag = diagnostic.Diagnostic("fatal",
|
|
"class body must contain only assignments and function definitions", {},
|
|
child.loc)
|
|
self.engine.process(diag)
|
|
|
|
if node.name in self.env_stack[-1]:
|
|
diag = diagnostic.Diagnostic("fatal",
|
|
"variable '{name}' is already defined", {"name":node.name}, node.name_loc)
|
|
self.engine.process(diag)
|
|
|
|
extractor = LocalExtractor(env_stack=self.env_stack, engine=self.engine)
|
|
extractor.visit(node)
|
|
|
|
# Now we create two types.
|
|
# The first type is the type of instances created by the constructor.
|
|
# Its attributes are those of the class environment, but wrapped
|
|
# appropriately so that they are linked to the class from which they
|
|
# originate.
|
|
instance_type = types.TInstance(node.name, OrderedDict())
|
|
|
|
# The second type is the type of the constructor itself (in other words,
|
|
# the class object): it is simply a singleton type that has the class
|
|
# environment as attributes.
|
|
constructor_type = types.TConstructor(instance_type)
|
|
constructor_type.attributes = extractor.typing_env
|
|
instance_type.constructor = constructor_type
|
|
|
|
self.env_stack[-1][node.name] = constructor_type
|
|
|
|
node = asttyped.ClassDefT(
|
|
constructor_type=constructor_type,
|
|
name=node.name,
|
|
bases=self.visit(node.bases), keywords=self.visit(node.keywords),
|
|
starargs=self.visit(node.starargs), kwargs=self.visit(node.kwargs),
|
|
body=node.body,
|
|
decorator_list=self.visit(node.decorator_list),
|
|
keyword_loc=node.keyword_loc, name_loc=node.name_loc,
|
|
lparen_loc=node.lparen_loc, star_loc=node.star_loc,
|
|
dstar_loc=node.dstar_loc, rparen_loc=node.rparen_loc,
|
|
colon_loc=node.colon_loc, at_locs=node.at_locs,
|
|
loc=node.loc)
|
|
|
|
try:
|
|
old_in_class, self.in_class = self.in_class, node
|
|
return self.generic_visit(node)
|
|
finally:
|
|
self.in_class = old_in_class
|
|
|
|
def visit_arg(self, node):
|
|
if node.annotation is not None:
|
|
diag = diagnostic.Diagnostic("fatal",
|
|
"type annotations are not supported here", {},
|
|
node.annotation.loc)
|
|
self.engine.process(diag)
|
|
|
|
return asttyped.argT(type=self._find_name(node.arg, node.loc),
|
|
arg=node.arg, annotation=None,
|
|
arg_loc=node.arg_loc, colon_loc=node.colon_loc, loc=node.loc)
|
|
|
|
def visit_Num(self, node):
|
|
if isinstance(node.n, int):
|
|
typ = builtins.TInt()
|
|
elif isinstance(node.n, float):
|
|
typ = builtins.TFloat()
|
|
else:
|
|
diag = diagnostic.Diagnostic("fatal",
|
|
"numeric type {type} is not supported", {"type": node.n.__class__.__name__},
|
|
node.loc)
|
|
self.engine.process(diag)
|
|
return asttyped.NumT(type=typ,
|
|
n=node.n, loc=node.loc)
|
|
|
|
def visit_Str(self, node):
|
|
if isinstance(node.s, str):
|
|
typ = builtins.TStr()
|
|
elif isinstance(node.s, bytes):
|
|
typ = builtins.TBytes()
|
|
else:
|
|
assert False
|
|
return asttyped.StrT(type=typ, s=node.s,
|
|
begin_loc=node.begin_loc, end_loc=node.end_loc, loc=node.loc)
|
|
|
|
def visit_Name(self, node):
|
|
return asttyped.NameT(type=self._find_name(node.id, node.loc),
|
|
id=node.id, ctx=node.ctx, loc=node.loc)
|
|
|
|
def visit_NameConstant(self, node):
|
|
if node.value is True or node.value is False:
|
|
typ = builtins.TBool()
|
|
elif node.value is None:
|
|
typ = builtins.TNone()
|
|
return asttyped.NameConstantT(type=typ, value=node.value, loc=node.loc)
|
|
|
|
def visit_Tuple(self, node):
|
|
node = self.generic_visit(node)
|
|
return asttyped.TupleT(type=types.TTuple([x.type for x in node.elts]),
|
|
elts=node.elts, ctx=node.ctx, loc=node.loc)
|
|
|
|
def visit_List(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.ListT(type=builtins.TList(),
|
|
elts=node.elts, ctx=node.ctx,
|
|
begin_loc=node.begin_loc, end_loc=node.end_loc, loc=node.loc)
|
|
return self.visit(node)
|
|
|
|
def visit_Attribute(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.AttributeT(type=types.TVar(),
|
|
value=node.value, attr=node.attr, ctx=node.ctx,
|
|
dot_loc=node.dot_loc, attr_loc=node.attr_loc, loc=node.loc)
|
|
return self.visit(node)
|
|
|
|
def visit_Slice(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.SliceT(type=types.TVar(),
|
|
lower=node.lower, upper=node.upper, step=node.step,
|
|
bound_colon_loc=node.bound_colon_loc,
|
|
step_colon_loc=node.step_colon_loc,
|
|
loc=node.loc)
|
|
return self.visit(node)
|
|
|
|
def visit_Subscript(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.SubscriptT(type=types.TVar(),
|
|
value=node.value, slice=node.slice, ctx=node.ctx,
|
|
begin_loc=node.begin_loc, end_loc=node.end_loc, loc=node.loc)
|
|
return self.visit(node)
|
|
|
|
def visit_BoolOp(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.BoolOpT(type=types.TVar(),
|
|
op=node.op, values=node.values,
|
|
op_locs=node.op_locs, loc=node.loc)
|
|
return self.visit(node)
|
|
|
|
def visit_UnaryOp(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.UnaryOpT(type=types.TVar(),
|
|
op=node.op, operand=node.operand,
|
|
loc=node.loc)
|
|
return self.visit(node)
|
|
|
|
def visit_BinOp(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.BinOpT(type=types.TVar(),
|
|
left=node.left, op=node.op, right=node.right,
|
|
loc=node.loc)
|
|
return self.visit(node)
|
|
|
|
def visit_Compare(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.CompareT(type=types.TVar(),
|
|
left=node.left, ops=node.ops, comparators=node.comparators,
|
|
loc=node.loc)
|
|
return self.visit(node)
|
|
|
|
def visit_IfExp(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.IfExpT(type=types.TVar(),
|
|
test=node.test, body=node.body, orelse=node.orelse,
|
|
if_loc=node.if_loc, else_loc=node.else_loc, loc=node.loc)
|
|
return self.visit(node)
|
|
|
|
def visit_ListComp(self, node):
|
|
extractor = LocalExtractor(env_stack=self.env_stack, engine=self.engine)
|
|
extractor.visit(node)
|
|
|
|
node = asttyped.ListCompT(
|
|
typing_env=extractor.typing_env, globals_in_scope=extractor.global_,
|
|
type=types.TVar(),
|
|
elt=node.elt, generators=node.generators,
|
|
begin_loc=node.begin_loc, end_loc=node.end_loc, loc=node.loc)
|
|
|
|
try:
|
|
self.env_stack.append(node.typing_env)
|
|
return self.generic_visit(node)
|
|
finally:
|
|
self.env_stack.pop()
|
|
|
|
def visit_Call(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.CallT(type=types.TVar(), iodelay=None, arg_exprs={},
|
|
remote_fn=False, func=node.func,
|
|
args=node.args, keywords=node.keywords,
|
|
starargs=node.starargs, kwargs=node.kwargs,
|
|
star_loc=node.star_loc, dstar_loc=node.dstar_loc,
|
|
begin_loc=node.begin_loc, end_loc=node.end_loc, loc=node.loc)
|
|
return node
|
|
|
|
def visit_Lambda(self, node):
|
|
extractor = LocalExtractor(env_stack=self.env_stack, engine=self.engine)
|
|
extractor.visit(node)
|
|
|
|
node = asttyped.LambdaT(
|
|
typing_env=extractor.typing_env, globals_in_scope=extractor.global_,
|
|
type=types.TVar(),
|
|
args=node.args, body=node.body,
|
|
lambda_loc=node.lambda_loc, colon_loc=node.colon_loc, loc=node.loc)
|
|
|
|
try:
|
|
self.env_stack.append(node.typing_env)
|
|
return self.generic_visit(node)
|
|
finally:
|
|
self.env_stack.pop()
|
|
|
|
def visit_ExceptHandler(self, node):
|
|
node = self.generic_visit(node)
|
|
if node.name is not None:
|
|
name_type = self._find_name(node.name, node.name_loc)
|
|
else:
|
|
name_type = types.TVar()
|
|
node = asttyped.ExceptHandlerT(
|
|
name_type=name_type,
|
|
filter=node.type, name=node.name, body=node.body,
|
|
except_loc=node.except_loc, as_loc=node.as_loc, name_loc=node.name_loc,
|
|
colon_loc=node.colon_loc, loc=node.loc)
|
|
return node
|
|
|
|
def visit_Raise(self, node):
|
|
node = self.generic_visit(node)
|
|
if node.cause:
|
|
diag = diagnostic.Diagnostic("error",
|
|
"'raise from' syntax is not supported", {},
|
|
node.from_loc)
|
|
self.engine.process(diag)
|
|
return node
|
|
|
|
def visit_For(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.ForT(
|
|
target=node.target, iter=node.iter, body=node.body, orelse=node.orelse,
|
|
trip_count=None, trip_interval=None,
|
|
keyword_loc=node.keyword_loc, in_loc=node.in_loc, for_colon_loc=node.for_colon_loc,
|
|
else_loc=node.else_loc, else_colon_loc=node.else_colon_loc, loc=node.loc)
|
|
return node
|
|
|
|
def visit_withitem(self, node):
|
|
node = self.generic_visit(node)
|
|
node = asttyped.withitemT(
|
|
context_expr=node.context_expr, optional_vars=node.optional_vars,
|
|
enter_type=types.TVar(), exit_type=types.TVar(),
|
|
as_loc=node.as_loc, loc=node.loc)
|
|
return node
|
|
|
|
# Unsupported visitors
|
|
#
|
|
def visit_unsupported(self, node):
|
|
diag = diagnostic.Diagnostic("fatal",
|
|
"this syntax is not supported", {},
|
|
node.loc)
|
|
self.engine.process(diag)
|
|
|
|
# expr
|
|
visit_Dict = visit_unsupported
|
|
visit_DictComp = visit_unsupported
|
|
visit_Ellipsis = visit_unsupported
|
|
visit_GeneratorExp = visit_unsupported
|
|
# visit_Set = visit_unsupported
|
|
visit_SetComp = visit_unsupported
|
|
visit_Starred = visit_unsupported
|
|
visit_Yield = visit_unsupported
|
|
visit_YieldFrom = visit_unsupported
|
|
|
|
# stmt
|
|
visit_Delete = visit_unsupported
|
|
visit_Import = visit_unsupported
|
|
visit_ImportFrom = visit_unsupported
|