Commits

Antonio Cuni  committed 86c2c43

factor out the logic for how to create an rpython wrapper from enforceargs, and put it in sourcetools

  • Participants
  • Parent commits c953de9
  • Branches autoreds

Comments (0)

Files changed (3)

File pypy/rlib/objectmodel.py

 import types
 import math
 import inspect
+from pypy.tool.sourcetools import rpython_wrapper
 
 # specialize is a decorator factory for attaching _annspecialcase_
 # attributes to functions: for example
                         f.func_name, srcargs[i], expected_type)
                     raise TypeError, msg
         #
-        # 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
+        template = """
+            def {name}({arglist}):
+                if not we_are_translated():
+                    typecheck({arglist})    # pypy.rlib.objectmodel
+                return {original}({arglist})
+        """
+        result = rpython_wrapper(f, template,
+                                 typecheck=typecheck,
+                                 we_are_translated=we_are_translated)
+        #
         srcargs, srcvarargs, srckeywords, defaults = inspect.getargspec(f)
         if kwds:
             types = tuple([kwds.get(arg) for arg in srcargs])
         assert len(srcargs) == len(types), (
             'not enough types provided: expected %d, got %d' %
             (len(types), len(srcargs)))
-        assert not srcvarargs, '*args not supported by enforceargs'
-        assert not srckeywords, '**kwargs not supported by enforceargs'
-        #
-        arglist = ', '.join(srcargs)
-        src = py.code.Source("""
-            def %(name)s(%(arglist)s):
-                if not we_are_translated():
-                    typecheck(%(arglist)s)    # pypy.rlib.objectmodel
-                return %(name)s_original(%(arglist)s)
-        """ % dict(name=f.func_name, arglist=arglist))
-        #
-        mydict = {f.func_name + '_original': f,
-                  'typecheck': typecheck,
-                  'we_are_translated': we_are_translated}
-        exec src.compile() in mydict
-        result = mydict[f.func_name]
-        result.func_defaults = f.func_defaults
-        result.func_dict.update(f.func_dict)
         result._annenforceargs_ = types
         return result
     return decorator
 
+
 # ____________________________________________________________
 
 class Symbolic(object):

File pypy/tool/sourcetools.py

     except AttributeError:
         firstlineno = -1
     return "(%s:%d)%s" % (mod or '?', firstlineno, name or 'UNKNOWN')
+
+
+def rpython_wrapper(f, template, **globaldict):
+    """  
+    We cannot simply wrap the function using *args, **kwds, because it's not
+    RPython. Instead, we generate a function from ``template`` with exactly
+    the same argument list.
+    """
+    srcargs, srcvarargs, srckeywords, defaults = inspect.getargspec(f)
+    assert not srcvarargs, '*args not supported by enforceargs'
+    assert not srckeywords, '**kwargs not supported by enforceargs'
+    #
+    arglist = ', '.join(srcargs)
+    src = template.format(name=f.func_name, arglist=arglist,
+                          original=f.func_name+'_original')
+    src = py.code.Source(src)
+    #
+    globaldict[f.func_name + '_original'] = f
+    exec src.compile() in globaldict
+    result = globaldict[f.func_name]
+    result.func_defaults = f.func_defaults
+    result.func_dict.update(f.func_dict)
+    return result

File pypy/tool/test/test_sourcetools.py

-from pypy.tool.sourcetools import func_with_new_name, func_renamer
+from pypy.tool.sourcetools import func_with_new_name, func_renamer, rpython_wrapper
 
 def test_rename():
     def f(x, y=5):
     bar3 = func_with_new_name(bar, 'bar3')
     assert bar3.func_doc == 'new doc'
     assert bar2.func_doc != bar3.func_doc
+
+
+def test_rpython_wrapper():
+    calls = []
+
+    def bar(a, b):
+        calls.append(('bar', a, b))
+        return a+b
+
+    template = """
+        def {name}({arglist}):
+            calls.append(('decorated', {arglist}))
+            return {original}({arglist})
+    """
+    bar = rpython_wrapper(bar, template, calls=calls)
+    assert bar(40, 2) == 42
+    assert calls == [
+        ('decorated', 40, 2),
+        ('bar', 40, 2),
+        ]
+
+