diff --git a/codetransformer/code.py b/codetransformer/code.py index c576206..6845b93 100644 --- a/codetransformer/code.py +++ b/codetransformer/code.py @@ -266,8 +266,9 @@ class Code: ---------- instrs : iterable of Instruction A sequence of codetransformer Instruction objects. - argnames : iterable of str, optional - The names of the arguments to the code object. + param_signature : iterable of str, optional + The names of the parameters to the code object. If a parameter name + starts with * or ** it will be interpreted as the vararg or varkeyword. name : str, optional The name of this code object. filename : str, optional @@ -298,6 +299,7 @@ class Code: lnotab name names + param_signature py_lnotab sparse_instrs stacksize @@ -306,6 +308,7 @@ class Code: __slots__ = ( '_instrs', '_argnames', + '_param_signature', '_argcount', '_kwonlyargcount', '_cellvars', @@ -320,7 +323,7 @@ class Code: def __init__(self, instrs, - argnames=(), + param_signature=(), *, cellvars=(), freevars=(), @@ -339,7 +342,7 @@ def __init__(self, _argnames = [] append_argname = _argnames.append varg = kwarg = None - for argname in argnames: + for argname in param_signature: if argname.startswith('**'): if kwarg is not None: raise ValueError('cannot specify **kwargs more than once') @@ -348,7 +351,9 @@ def __init__(self, elif argname.startswith('*'): if varg is not None: raise ValueError('cannot specify *args more than once') - varg = argname[1:] + if argname != '*': + # lone star, not varg + varg = argname[1:] argcounter = kwonlyargcount # all following args are kwonly. continue argcounter[0] += 1 @@ -376,6 +381,7 @@ def __init__(self, self._instrs = instrs self._argnames = tuple(_argnames) + self._param_signature = tuple(param_signature) self._argcount = argcount[0] self._kwonlyargcount = kwonlyargcount[0] self._cellvars = cellvars @@ -497,7 +503,7 @@ def from_pycode(cls, co): return cls( filter(bool, sparse_instrs), - argnames=new_paramnames, + param_signature=new_paramnames, cellvars=co.co_cellvars, freevars=co.co_freevars, name=co.co_name, @@ -618,7 +624,7 @@ def sparse_instrs(self): def argcount(self): """The number of arguments this code object accepts. - This does not include varargs (\*args). + This does not include varargs (*args). """ return self._argcount @@ -626,7 +632,7 @@ def argcount(self): def kwonlyargcount(self): """The number of keyword only arguments this code object accepts. - This does not include varkwargs (\*\*kwargs). + This does not include varkwargs (**kwargs). """ return self._kwonlyargcount @@ -665,6 +671,15 @@ def argnames(self): """ return self._argnames + @property + def param_signature(self): + """The string signature for all the parameters to this code object. + + Each parameter name is a string. The vararg will be prefixed with '*' + and the varkeyword will be prefixed with '**'. + """ + return self._param_signature + @property def varnames(self): """The names of all of the local variables in this code object. diff --git a/codetransformer/core.py b/codetransformer/core.py index 4876014..63060d4 100644 --- a/codetransformer/core.py +++ b/codetransformer/core.py @@ -193,7 +193,7 @@ def transform(self, code, *, name=None, filename=None): return Code( post_transform, - code.argnames, + code.param_signature, cellvars=self.transform_cellvars(code.cellvars), freevars=self.transform_freevars(code.freevars), name=name if name is not None else code.name, diff --git a/codetransformer/tests/test_code.py b/codetransformer/tests/test_code.py index f9135e8..8ec472e 100644 --- a/codetransformer/tests/test_code.py +++ b/codetransformer/tests/test_code.py @@ -7,7 +7,12 @@ import pytest from codetransformer.code import Code, Flag, pycode -from codetransformer.instructions import LOAD_CONST, LOAD_FAST, uses_free +from codetransformer.instructions import ( + LOAD_CONST, + LOAD_FAST, + RETURN_VALUE, + uses_free, +) @pytest.fixture(scope='module') @@ -178,7 +183,7 @@ def abc_code(): a = LOAD_CONST('a') b = LOAD_CONST('b') c = LOAD_CONST('c') # not in instrs - code = Code((a, b), argnames=()) + code = Code((a, b), param_signature=()) return (a, b, c), code @@ -220,3 +225,26 @@ def code(): # pragma: no cover buf = StringIO() code.dis(file=buf) assert buf.getvalue() == expected + + +@pytest.mark.parametrize( + 'param_signature, has_varargs, has_varkeywords', ( + (('a',), False, False), + (('a', 'b'), False, False), + (('a', 'b', '*c'), True, False), + (('a', 'b', '**c'), False, True), + (('a', 'b', '*', 'c'), False, False), + (('a', 'b', '*c', 'd'), True, False), + (('a', 'b', '*c', '**d'), True, True), + (('a', 'b', '*c', 'd', '**e'), True, True), + ), +) +def test_param_signature(param_signature, has_varargs, has_varkeywords): + code = Code( + (LOAD_CONST(None), RETURN_VALUE()), + param_signature=param_signature, + ) + assert code.param_signature == param_signature + + assert code.flags['CO_VARARGS'] == has_varargs + assert code.flags['CO_VARKEYWORDS'] == has_varkeywords diff --git a/codetransformer/tests/test_core.py b/codetransformer/tests/test_core.py index 1920ff0..1601bbe 100644 --- a/codetransformer/tests/test_core.py +++ b/codetransformer/tests/test_core.py @@ -128,3 +128,49 @@ class c(CodeTransformer): c.context assert str(e.value) == 'no active transformation context' + + +def test_nop(): + result = object() + + @CodeTransformer() + def f(): + return result + + assert f() is result + + @CodeTransformer() + def f(a): + return a + + assert f(result) is result + + @CodeTransformer() + def f(a=result): + return a + + assert f() is result + + @CodeTransformer() + def f(*args): + return args[0] + + assert f(result) is result + + @CodeTransformer() + def f(*args, a): + return a + + assert f(a=result) is result + + @CodeTransformer() + def f(**kwargs): + return kwargs['a'] + + assert f(a=result) is result + + @CodeTransformer() + def f(a, *b, c, **d): + return a, b, c, d + + assert f(1, 2, c=3, d=4) == (1, (2,), 3, {'d': 4})