PEP 3107 - Function Annotations thanks to Tony Lownds
diff --git a/Lib/compiler/ast.py b/Lib/compiler/ast.py
index d21905f..ac48710 100644
--- a/Lib/compiler/ast.py
+++ b/Lib/compiler/ast.py
@@ -33,7 +33,10 @@
pass # implemented by subclasses
class EmptyNode(Node):
- pass
+ def getChildNodes(self):
+ return ()
+ def getChildren(self):
+ return ()
class Expression(Node):
# Expression is an artificial node class to support "eval"
@@ -487,12 +490,13 @@
return "From(%s, %s, %s)" % (repr(self.modname), repr(self.names), repr(self.level))
class Function(Node):
- def __init__(self, decorators, name, argnames, defaults, kwonlyargs, flags, doc, code, lineno=None):
+ def __init__(self, decorators, name, arguments, defaults, kwonlyargs, returns, flags, doc, code, lineno=None):
self.decorators = decorators
self.name = name
- self.argnames = argnames
+ self.arguments = arguments
self.defaults = defaults
self.kwonlyargs = kwonlyargs
+ self.returns = returns
self.flags = flags
self.doc = doc
self.code = code
@@ -508,9 +512,10 @@
children = []
children.append(self.decorators)
children.append(self.name)
- children.append(self.argnames)
+ children.extend(flatten(self.arguments))
children.extend(flatten(self.defaults))
- children.append(self.kwonlyargs)
+ children.extend(flatten(self.kwonlyargs))
+ children.append(self.returns)
children.append(self.flags)
children.append(self.doc)
children.append(self.code)
@@ -520,18 +525,22 @@
nodelist = []
if self.decorators is not None:
nodelist.append(self.decorators)
+ nodelist.extend(flatten_nodes(self.arguments))
nodelist.extend(flatten_nodes(self.defaults))
+ nodelist.extend(flatten_nodes(self.kwonlyargs))
+ if self.returns is not None:
+ nodelist.append(self.returns)
nodelist.append(self.code)
return tuple(nodelist)
def __repr__(self):
- return "Function(%s, %s, %s, %s, %s, %s, %s, %s)" % (repr(self.decorators), repr(self.name), repr(self.argnames), repr(self.defaults), repr(self.kwonlyargs), repr(self.flags), repr(self.doc), repr(self.code))
+ return "Function(%s, %s, %s, %s, %s, %s, %s, %s, %s)" % (repr(self.decorators), repr(self.name), repr(self.arguments), repr(self.defaults), repr(self.kwonlyargs), repr(self.returns), repr(self.flags), repr(self.doc), repr(self.code))
class GenExpr(Node):
def __init__(self, code, lineno=None):
self.code = code
self.lineno = lineno
- self.argnames = ['.0']
+ self.arguments = [SimpleArg('.0', None)]
self.varargs = self.kwargs = None
self.kwonlyargs = ()
@@ -715,9 +724,24 @@
def __repr__(self):
return "Keyword(%s, %s)" % (repr(self.name), repr(self.expr))
+class Kwarg(Node):
+ def __init__(self, arg, expr, lineno=None):
+ self.arg = arg
+ self.expr = expr
+ self.lineno = lineno
+
+ def getChildren(self):
+ return self.arg, self.expr
+
+ def getChildNodes(self):
+ return self.arg, self.expr
+
+ def __repr__(self):
+ return "Kwarg(%s, %s)" % (repr(self.arg), repr(self.expr))
+
class Lambda(Node):
- def __init__(self, argnames, defaults, kwonlyargs, flags, code, lineno=None):
- self.argnames = argnames
+ def __init__(self, arguments, defaults, kwonlyargs, flags, code, lineno=None):
+ self.arguments = arguments
self.defaults = defaults
self.kwonlyargs = kwonlyargs
self.flags = flags
@@ -728,25 +752,28 @@
self.varargs = 1
if flags & CO_VARKEYWORDS:
self.kwargs = 1
+ self.returns = None
def getChildren(self):
children = []
- children.append(self.argnames)
+ children.extend(flatten(self.arguments))
children.extend(flatten(self.defaults))
- children.append(self.kwonlyargs)
+ children.extend(flatten(self.kwonlyargs))
children.append(self.flags)
children.append(self.code)
return tuple(children)
def getChildNodes(self):
nodelist = []
+ nodelist.extend(flatten_nodes(self.arguments))
nodelist.extend(flatten_nodes(self.defaults))
+ nodelist.extend(flatten_nodes(self.kwonlyargs))
nodelist.append(self.code)
return tuple(nodelist)
def __repr__(self):
- return "Lambda(%s, %s, %s, %s, %s)" % (repr(self.argnames), repr(self.defaults), repr(self.kwonlyargs), repr(self.flags), repr(self.code))
+ return "Lambda(%s, %s, %s, %s, %s)" % (repr(self.arguments), repr(self.defaults), repr(self.kwonlyargs), repr(self.flags), repr(self.code))
class LeftShift(Node):
def __init__(self, (left, right), lineno=None):
@@ -897,6 +924,22 @@
def __repr__(self):
return "Name(%s)" % (repr(self.name),)
+class NestedArgs(Node):
+ def __init__(self, args, lineno=None):
+ self.args = args
+ self.lineno = lineno
+
+ def getChildren(self):
+ return tuple(flatten(self.args))
+
+ def getChildNodes(self):
+ nodelist = []
+ nodelist.extend(flatten_nodes(self.args))
+ return tuple(nodelist)
+
+ def __repr__(self):
+ return "NestedArgs(%s)" % (repr(self.args),)
+
class Not(Node):
def __init__(self, expr, lineno=None):
self.expr = expr
@@ -1071,6 +1114,27 @@
def __repr__(self):
return "Set(%s)" % (repr(self.items),)
+class SimpleArg(Node):
+ def __init__(self, name, annotation, lineno=None):
+ self.name = name
+ self.annotation = annotation
+ self.lineno = lineno
+
+ def getChildren(self):
+ children = []
+ children.append(self.name)
+ children.append(self.annotation)
+ return tuple(children)
+
+ def getChildNodes(self):
+ nodelist = []
+ if self.annotation is not None:
+ nodelist.append(self.annotation)
+ return tuple(nodelist)
+
+ def __repr__(self):
+ return "SimpleArg(%s, %s)" % (repr(self.name), repr(self.annotation))
+
class Slice(Node):
def __init__(self, expr, flags, lower, upper, lineno=None):
self.expr = expr
diff --git a/Lib/compiler/pyassem.py b/Lib/compiler/pyassem.py
index 7f44d46..86e2c49 100644
--- a/Lib/compiler/pyassem.py
+++ b/Lib/compiler/pyassem.py
@@ -314,7 +314,7 @@
super_init = FlowGraph.__init__
def __init__(self, name, filename,
- args=(), kwonlyargs={}, optimized=0, klass=None):
+ args=(), kwonlyargs=(), optimized=0, klass=None):
self.super_init()
self.name = name
self.filename = filename
@@ -338,24 +338,40 @@
# The offsets used by LOAD_CLOSURE/LOAD_DEREF refer to both
# kinds of variables.
self.closure = []
- self.varnames = list(args) or []
- for i in range(len(self.varnames)):
+ # The varnames list needs to be computed after flags have been set
+ self.varnames = []
+ self.stage = RAW
+
+ def computeVarnames(self):
+ # self.args is positional, vararg, kwarg, kwonly, unpacked. This
+ # order is due to the visit order in symbol module and could change.
+ # argcount is # len(self.args) - len(unpacked). We want
+ # self.varnames to be positional, kwonly, vararg, kwarg, unpacked
+ # and argcount to be len(positional).
+
+ # determine starting index of unpacked, kwonly, vararg
+ u = self.argcount # starting index of unpacked
+ k = u - len(self.kwonlyargs) # starting index of kwonly
+ v = k - self.checkFlag(CO_VARARGS) - self.checkFlag(CO_VARKEYWORDS)
+
+ vars = list(self.args)
+ self.varnames = vars[:v] + vars[k:u] + vars[v:k] + vars[u:]
+ self.argcount = v
+
+ # replace TupleArgs with calculated var name
+ for i in range(self.argcount):
var = self.varnames[i]
if isinstance(var, TupleArg):
self.varnames[i] = var.getName()
- self.stage = RAW
def setDocstring(self, doc):
self.docstring = doc
def setFlag(self, flag):
self.flags = self.flags | flag
- if flag == CO_VARARGS:
- self.argcount = self.argcount - 1
def checkFlag(self, flag):
- if self.flags & flag:
- return 1
+ return (self.flags & flag) == flag
def setFreeVars(self, names):
self.freevars = list(names)
@@ -366,6 +382,7 @@
def getCode(self):
"""Get a Python code object"""
assert self.stage == RAW
+ self.computeVarnames()
self.computeStackDepth()
self.flattenGraph()
assert self.stage == FLAT
@@ -575,6 +592,12 @@
lnotab.nextLine(oparg)
continue
hi, lo = twobyte(oparg)
+
+ extended, hi = twobyte(hi)
+ if extended:
+ ehi, elo = twobyte(extended)
+ lnotab.addCode(self.opnum['EXTENDED_ARG'], elo, ehi)
+
try:
lnotab.addCode(self.opnum[opname], lo, hi)
except ValueError:
@@ -595,8 +618,6 @@
else:
nlocals = len(self.varnames)
argcount = self.argcount
- if self.flags & CO_VARKEYWORDS:
- argcount = argcount - 1
kwonlyargcount = len(self.kwonlyargs)
return new.code(argcount, kwonlyargcount,
nlocals, self.stacksize, self.flags,
@@ -809,7 +830,8 @@
return self.CALL_FUNCTION(argc)-2
def MAKE_FUNCTION(self, argc):
hi, lo = divmod(argc, 256)
- return -(lo + hi * 2)
+ ehi, hi = divmod(hi, 256)
+ return -(lo + hi * 2 + ehi)
def MAKE_CLOSURE(self, argc):
# XXX need to account for free variables too!
return -argc
diff --git a/Lib/compiler/pycodegen.py b/Lib/compiler/pycodegen.py
index 353c2c9..325ca06 100644
--- a/Lib/compiler/pycodegen.py
+++ b/Lib/compiler/pycodegen.py
@@ -378,18 +378,57 @@
walk(node.code, gen)
gen.finish()
self.set_lineno(node)
+ num_kwargs = 0
for keyword in node.kwonlyargs:
default = keyword.expr
if isinstance(default, ast.EmptyNode):
continue
- self.emit('LOAD_CONST', keyword.name)
+ self.emit('LOAD_CONST', keyword.arg.name)
self.visit(default)
+ num_kwargs += 1
for default in node.defaults:
self.visit(default)
- self._makeClosure(gen, len(node.defaults))
+
+ num_annotations = self._visit_annotations(node)
+
+ oparg = len(node.defaults)
+ oparg |= num_kwargs << 8
+ oparg |= num_annotations << 16
+
+ self._makeClosure(gen, oparg)
for i in range(ndecorators):
self.emit('CALL_FUNCTION', 1)
+ def _visit_annotations(self, node):
+ # emit code, return num_annotations
+ annotations = []
+ annotations.extend(self._visit_argument_annotations(node.arguments))
+ annotations.extend(self._visit_kwarg_annotations(node.kwonlyargs))
+ if node.returns:
+ self.visit(node.returns)
+ annotations.append('return')
+ if not annotations:
+ return 0
+ self.emit('LOAD_CONST', tuple(annotations))
+ return len(annotations) + 1
+
+ def _visit_argument_annotations(self, arguments):
+ for arg in arguments:
+ if isinstance(arg, ast.SimpleArg):
+ if arg.annotation:
+ self.visit(arg.annotation)
+ yield arg.name
+ else:
+ for name in self._visit_argument_annotations(arg.args):
+ yield name
+
+ def _visit_kwarg_annotations(self, kwargs):
+ for kwarg in kwargs:
+ arg = kwarg.arg
+ if arg.annotation:
+ self.visit(arg.annotation)
+ yield arg.name
+
def visitClass(self, node):
gen = self.ClassGen(node, self.scopes,
self.get_module())
@@ -1323,7 +1362,7 @@
else:
name = func.name
- args, hasTupleArg = generateArgList(func.argnames)
+ args, hasTupleArg = generateArgList(func.arguments)
kwonlyargs = generateKwonlyArgList(func.kwonlyargs)
self.graph = pyassem.PyFlowGraph(name, func.filename, args,
kwonlyargs=kwonlyargs,
@@ -1334,7 +1373,7 @@
if not isLambda and func.doc:
self.setDocstring(func.doc)
- lnf = walk(func.code, self.NameFinder(args), verbose=0)
+ lnf = walk(func.code, self.NameFinder(args+kwonlyargs), verbose=0)
self.locals.push(lnf.getLocals())
if func.varargs:
self.graph.setFlag(CO_VARARGS)
@@ -1342,7 +1381,7 @@
self.graph.setFlag(CO_VARKEYWORDS)
self.set_lineno(func)
if hasTupleArg:
- self.generateArgUnpack(func.argnames)
+ self.generateArgUnpack(func.arguments)
def get_module(self):
return self.module
@@ -1356,9 +1395,9 @@
def generateArgUnpack(self, args):
for i in range(len(args)):
arg = args[i]
- if isinstance(arg, tuple):
+ if isinstance(arg, ast.NestedArgs):
self.emit('LOAD_FAST', '.%d' % (i * 2))
- self.unpackSequence(arg)
+ self.unpackSequence(tuple(_nested_names(arg)))
def unpackSequence(self, tup):
if VERSION > 1:
@@ -1452,21 +1491,29 @@
count = 0
for i in range(len(arglist)):
elt = arglist[i]
- if isinstance(elt, str):
- args.append(elt)
- elif isinstance(elt, tuple):
- args.append(TupleArg(i * 2, elt))
- extra.extend(misc.flatten(elt))
+ if isinstance(elt, ast.SimpleArg):
+ args.append(elt.name)
+ elif isinstance(elt, ast.NestedArgs):
+ t = tuple(_nested_names(elt))
+ args.append(TupleArg(i * 2, t))
+ extra.extend(misc.flatten(t))
count = count + 1
else:
raise ValueError, "unexpect argument type:", elt
return args + extra, count
+def _nested_names(elt):
+ for arg in elt.args:
+ if isinstance(arg, ast.SimpleArg):
+ yield arg.name
+ elif isinstance(arg, ast.NestedArgs):
+ yield tuple(_nested_names(arg))
+
def generateKwonlyArgList(keywordOnlyArgs):
- kwonlyargs = {}
+ kwonlyargs = []
for elt in keywordOnlyArgs:
- assert isinstance(elt, ast.Keyword)
- kwonlyargs[elt.name] = elt.expr
+ assert isinstance(elt, ast.Kwarg)
+ kwonlyargs.append(elt.arg.name)
return kwonlyargs
def findOp(node):
diff --git a/Lib/compiler/symbols.py b/Lib/compiler/symbols.py
index 7ddf42c..3585efc 100644
--- a/Lib/compiler/symbols.py
+++ b/Lib/compiler/symbols.py
@@ -233,7 +233,12 @@
if parent.nested or isinstance(parent, FunctionScope):
scope.nested = 1
self.scopes[node] = scope
- self._do_args(scope, node.argnames)
+
+ args = node.arguments
+ for kwonly in node.kwonlyargs:
+ args.append(kwonly.arg)
+ self._do_arguments(scope, args)
+
self.visit(node.code, scope)
self.handle_free_vars(scope, parent)
@@ -275,16 +280,18 @@
if parent.nested or isinstance(parent, FunctionScope):
scope.nested = 1
self.scopes[node] = scope
- self._do_args(scope, node.argnames)
+ self._do_arguments(scope, node.arguments)
self.visit(node.code, scope)
self.handle_free_vars(scope, parent)
- def _do_args(self, scope, args):
- for name in args:
- if type(name) == types.TupleType:
- self._do_args(scope, name)
+ def _do_arguments(self, scope, arguments):
+ for node in arguments:
+ if isinstance(node, ast.SimpleArg):
+ scope.add_param(node.name)
+ if node.annotation:
+ self.visit(node.annotation, scope)
else:
- scope.add_param(name)
+ self._do_arguments(scope, node.args)
def handle_free_vars(self, scope, parent):
parent.add_child(scope)
diff --git a/Lib/compiler/transformer.py b/Lib/compiler/transformer.py
index dc88222..4a8e623 100644
--- a/Lib/compiler/transformer.py
+++ b/Lib/compiler/transformer.py
@@ -234,25 +234,24 @@
return Decorators(items)
def funcdef(self, nodelist):
- # -6 -5 -4 -3 -2 -1
- # funcdef: [decorators] 'def' NAME parameters ':' suite
- # parameters: '(' [varargslist] ')'
-
- if len(nodelist) == 6:
- assert nodelist[0][0] == symbol.decorators
+ # 0 1 2 4 -1
+ # funcdef: [decorators] 'def' NAME parameters ['->' test] ':' suite
+ # parameters: '(' [typedargslist] ')'
+ if nodelist[0][0] == symbol.decorators:
decorators = self.decorators(nodelist[0][1:])
+ nodelist = nodelist[1:]
else:
- assert len(nodelist) == 5
decorators = None
+ assert len(nodelist) in (5, 7)
- lineno = nodelist[-4][2]
- name = nodelist[-4][1]
- args = nodelist[-3][2]
+ lineno = nodelist[0][2]
+ name = nodelist[1][1]
+ args = nodelist[2][2]
- if args[0] == symbol.varargslist:
- names, defaults, kwonlyargs, flags = self.com_arglist(args[1:])
+ if args[0] == symbol.varargslist or args[0] == symbol.typedargslist:
+ arguments, defaults, kwonly, flags = self.com_arglist(args[1:])
else:
- names = defaults = kwonlyargs = ()
+ arguments = defaults = kwonly = ()
flags = 0
doc = self.get_docstring(nodelist[-1])
@@ -263,22 +262,28 @@
assert isinstance(code, Stmt)
assert isinstance(code.nodes[0], Discard)
del code.nodes[0]
- return Function(decorators, name, names, defaults,
- kwonlyargs, flags, doc, code, lineno=lineno)
+
+ if len(nodelist) == 7:
+ returns = self.com_node(nodelist[4])
+ else:
+ returns = None
+
+ return Function(decorators, name, arguments, defaults,
+ kwonly, returns, flags, doc, code, lineno=lineno)
def lambdef(self, nodelist):
# lambdef: 'lambda' [varargslist] ':' test
if nodelist[2][0] == symbol.varargslist:
- names, defaults, kwonlyargs, flags = \
+ arguments, defaults, kwonlyargs, flags = \
self.com_arglist(nodelist[2][1:])
else:
- names = defaults = kwonlyargs = ()
+ arguments = defaults = kwonlyargs = ()
flags = 0
# code for lambda
code = self.com_node(nodelist[-1])
- return Lambda(names, defaults, kwonlyargs,
+ return Lambda(arguments, defaults, kwonlyargs,
flags, code, lineno=nodelist[1][2])
old_lambdef = lambdef
@@ -324,10 +329,25 @@
def varargslist(self, nodelist):
raise WalkerError
- def fpdef(self, nodelist):
+ def vfpdef(self, nodelist):
raise WalkerError
- def fplist(self, nodelist):
+ def vfplist(self, nodelist):
+ raise WalkerError
+
+ def vname(self, nodelist):
+ raise WalkerError
+
+ def typedargslist(self, nodelist):
+ raise WalkerError
+
+ def tfpdef(self, nodelist):
+ raise WalkerError
+
+ def tfplist(self, nodelist):
+ raise WalkerError
+
+ def tname(self, nodelist):
raise WalkerError
def dotted_name(self, nodelist):
@@ -786,9 +806,10 @@
return Discard(Const(None))
def keywordonlyargs(self, nodelist):
- # (',' NAME ['=' test])*
+ # (',' tname ['=' test])*
# ^^^
# ------+
+ # tname and vname are handled.
kwonlyargs = []
i = 0
while i < len(nodelist):
@@ -802,10 +823,25 @@
i += 2
if node[0] == token.DOUBLESTAR:
return kwonlyargs, i
- elif node[0] == token.NAME:
- kwonlyargs.append(Keyword(node[1], default, lineno=node[2]))
+ elif node[0] in (symbol.vname, symbol.tname):
+ lineno = extractLineNo(node)
+ kwarg = Kwarg(self._simplearg(node), default, lineno=lineno)
+ kwonlyargs.append(kwarg)
i += 2
return kwonlyargs, i
+
+ def _simplearg(self, node):
+ # tname: NAME [':' test]
+ # vname: NAME
+ assert node[0] == symbol.vname or node[0] == symbol.tname
+ name = node[1][1]
+ lineno = node[1][2]
+ assert isinstance(name, str)
+ if len(node) > 2:
+ annotation = self.com_node(node[3])
+ else:
+ annotation = None
+ return SimpleArg(name, annotation, lineno)
def com_arglist(self, nodelist):
# varargslist:
@@ -814,7 +850,7 @@
# | fpdef ['=' test] (',' fpdef ['=' test])* [',']
# fpdef: NAME | '(' fplist ')'
# fplist: fpdef (',' fpdef)* [',']
- names = []
+ arguments = []
kwonlyargs = []
defaults = []
flags = 0
@@ -825,14 +861,15 @@
if node[0] == token.STAR or node[0] == token.DOUBLESTAR:
if node[0] == token.STAR:
node = nodelist[i+1]
- if node[0] == token.NAME: # vararg
- names.append(node[1])
+ if node[0] in (symbol.tname, symbol.vname): # vararg
+ arguments.append(self._simplearg(node))
flags = flags | CO_VARARGS
i = i + 3
else: # no vararg
assert node[0] == token.COMMA
i += 2
- if i < len(nodelist) and nodelist[i][0] == token.NAME:
+ if i < len(nodelist) and \
+ nodelist[i][0] in (symbol.tname, symbol.vname):
kwonlyargs, skip = self.keywordonlyargs(nodelist[i:])
i += skip
@@ -843,13 +880,13 @@
node = nodelist[i+1]
else:
raise ValueError, "unexpected token: %s" % t
- names.append(node[1])
+ arguments.append(self._simplearg(node))
flags = flags | CO_VARKEYWORDS
break
- # fpdef: NAME | '(' fplist ')'
- names.append(self.com_fpdef(node))
+ # tfpdef: tname | '(' tfplist ')'
+ arguments.append(self.com_tfpdef(node))
i = i + 1
if i < len(nodelist) and nodelist[i][0] == token.EQUAL:
@@ -863,21 +900,24 @@
# skip the comma
i = i + 1
- return names, defaults, kwonlyargs, flags
+ return arguments, defaults, kwonlyargs, flags
- def com_fpdef(self, node):
- # fpdef: NAME | '(' fplist ')'
+ def com_tfpdef(self, node):
+ # tfpdef: tname | '(' tfplist ')'
+ # def f((x)): -- x is not nested
+ while node[1][0] == token.LPAR and len(node[2]) == 2:
+ node = node[2][1]
if node[1][0] == token.LPAR:
- return self.com_fplist(node[2])
- return node[1][1]
+ return NestedArgs(self.com_tfplist(node[2]))
+ return self._simplearg(node[1])
- def com_fplist(self, node):
- # fplist: fpdef (',' fpdef)* [',']
+ def com_tfplist(self, node):
+ # tfplist: tfpdef (',' tfpdef)* [',']
if len(node) == 2:
- return self.com_fpdef(node[1])
+ return self.com_tfpdef(node[1]),
list = []
for i in range(1, len(node), 2):
- list.append(self.com_fpdef(node[i]))
+ list.append(self.com_tfpdef(node[i]))
return tuple(list)
def com_dotted_name(self, node):