[pypy-commit] pypy better-enforceargs: do some typechecking on @enforceargs functions

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: better-enforceargs
Changeset: r56106:0214d225faf9
Date: 2012-07-17 17:52 +0200
http://bitbucket.org/pypy/pypy/changeset/0214d225faf9/

Log:	do some typechecking on @enforceargs functions

diff --git a/pypy/rlib/objectmodel.py b/pypy/rlib/objectmodel.py
--- a/pypy/rlib/objectmodel.py
+++ b/pypy/rlib/objectmodel.py
@@ -3,6 +3,7 @@
 RPython-compliant way.
 """
 
+import py
 import sys
 import types
 import math
@@ -106,15 +107,43 @@
 
 specialize = _Specialize()
 
-def enforceargs(*args):
+def enforceargs(*types):
     """ Decorate a function with forcing of RPython-level types on arguments.
     None means no enforcing.
 
     XXX shouldn't we also add asserts in function body?
     """
+    import inspect
     def decorator(f):
-        f._annenforceargs_ = args
-        return f
+        def typecheck(*args):
+            for t, arg in zip(types, args):
+                if t is not None and not isinstance(arg, t):
+                    raise TypeError
+        #
+        # we cannot simply wrap the function using *args, **kwds, because it's
+        # not RPython. Instead, we generate a function with exactly the same
+        # argument list
+        argspec = inspect.getargspec(f)
+        assert len(argspec.args) == len(types), (
+            'not enough types provided: expected %d, got %d' %
+            (len(types), len(argspec.args)))
+        assert not argspec.varargs, '*args not supported by enforceargs'
+        assert not argspec.keywords, '**kwargs not supported by enforceargs'
+        #
+        arglist = ', '.join(argspec.args)
+        src = py.code.Source("""
+            def {name}({arglist}):
+                typecheck({arglist})
+                return {name}_original({arglist})
+        """.format(name=f.func_name, arglist=arglist))
+        #
+        mydict = {f.func_name + '_original': f,
+                  'typecheck': typecheck}
+        exec src.compile() in mydict
+        result = mydict[f.func_name]
+        # XXX defaults
+        result._annenforceargs_ = types
+        return result
     return decorator
 
 # ____________________________________________________________
diff --git a/pypy/rlib/test/test_objectmodel.py b/pypy/rlib/test/test_objectmodel.py
--- a/pypy/rlib/test/test_objectmodel.py
+++ b/pypy/rlib/test/test_objectmodel.py
@@ -420,9 +420,14 @@
 def test_enforceargs_decorator():
     @enforceargs(int, str, None)
     def f(a, b, c):
-        pass
+        return a, b, c
+    assert f._annenforceargs_ == (int, str, None)
+    assert f.func_name == 'f'
+    assert f(1, 'hello', 42) == (1, 'hello', 42)
+    py.test.raises(TypeError, "f(1, 2, 3)")
+    py.test.raises(TypeError, "f('hello', 'world', 3)")
 
-    assert f._annenforceargs_ == (int, str, None)
+    
 
 def getgraph(f, argtypes):
     from pypy.translator.translator import TranslationContext, graphof