Skip to content

Commit

Permalink
compiler.embedding: implement type annotations for function arguments.
Browse files Browse the repository at this point in the history
Fixes #318.
whitequark committed Aug 8, 2016
1 parent 8a243d3 commit 5a2306a
Showing 3 changed files with 69 additions and 1 deletion.
47 changes: 47 additions & 0 deletions artiq/compiler/embedding.py
Original file line number Diff line number Diff line change
@@ -415,6 +415,53 @@ def __init__(self, engine, prelude, globals, host_environment, quote):
self.host_environment = host_environment
self.quote = quote

def match_annotation(self, annot):
if isinstance(annot, ast.Name):
if annot.id == "TNone":
return builtins.TNone()
if annot.id == "TBool":
return builtins.TBool()
if annot.id == "TInt32":
return builtins.TInt(types.TValue(32))
if annot.id == "TInt64":
return builtins.TInt(types.TValue(64))
if annot.id == "TFloat":
return builtins.TFloat()
if annot.id == "TStr":
return builtins.TStr()
if annot.id == "TRange32":
return builtins.TRange(builtins.TInt(types.TValue(32)))
if annot.id == "TRange64":
return builtins.TRange(builtins.TInt(types.TValue(64)))
if annot.id == "TVar":
return types.TVar()
elif (isinstance(annot, ast.Call) and
annot.keywords is None and
annot.starargs is None and
annot.kwargs is None and
isinstance(annot.func, ast.Name)):
if annot.func.id == "TList" and len(annot.args) == 1:
elttyp = self.match_annotation(annot.args[0])
if elttyp is not None:
return builtins.TList()
else:
return None

if annot is not None:
diag = diagnostic.Diagnostic("error",
"unrecognized type annotation", {},
annot.loc)
self.engine.process(diag)

def visit_arg(self, node):
typ = self._find_name(node.arg, node.loc)
annottyp = self.match_annotation(node.annotation)
if annottyp is not None:
typ.unify(annottyp)
return asttyped.argT(type=typ,
arg=node.arg, annotation=None,
arg_loc=node.arg_loc, colon_loc=node.colon_loc, loc=node.loc)

def visit_quoted_function(self, node, function):
extractor = LocalExtractor(env_stack=self.env_stack, engine=self.engine)
extractor.visit(node)
8 changes: 7 additions & 1 deletion artiq/compiler/transforms/asttyped_rewriter.py
Original file line number Diff line number Diff line change
@@ -307,8 +307,14 @@ def visit_ClassDef(self, node):
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=self.visit(node.annotation),
arg=node.arg, annotation=None,
arg_loc=node.arg_loc, colon_loc=node.colon_loc, loc=node.loc)

def visit_Num(self, node):
15 changes: 15 additions & 0 deletions artiq/test/coredevice/test_embedding.py
Original file line number Diff line number Diff line change
@@ -137,6 +137,21 @@ def test_args(self):
exp.builtin()


class _Annotation(EnvExperiment):
def build(self):
self.setattr_device("core")

@kernel
def overflow(self, x: TInt64) -> TBool:
return (x << 32) != 0


class AnnotationTest(ExperimentCase):
def test_annotation(self):
exp = self.create(_Annotation)
self.assertEqual(exp.overflow(1), True)


class _Payload1MB(EnvExperiment):
def build(self):
self.setattr_device("core")

0 comments on commit 5a2306a

Please sign in to comment.