[pypy-svn] r71341 - in pypy/branch/string-promote/pypy/rlib: . test

fijal at codespeak.net fijal at codespeak.net
Fri Feb 19 20:40:20 CET 2010


Author: fijal
Date: Fri Feb 19 20:40:18 2010
New Revision: 71341

Modified:
   pypy/branch/string-promote/pypy/rlib/jit.py
   pypy/branch/string-promote/pypy/rlib/test/test_jit.py
Log:
Grow function_promote an optional arg and a test


Modified: pypy/branch/string-promote/pypy/rlib/jit.py
==============================================================================
--- pypy/branch/string-promote/pypy/rlib/jit.py	(original)
+++ pypy/branch/string-promote/pypy/rlib/jit.py	Fri Feb 19 20:40:18 2010
@@ -25,25 +25,28 @@
     func._jit_loop_invariant_ = True
     return func
 
-
-def purefunction_promote(func):
-    import inspect
-    purefunction(func)
-    func._pure_function_with_all_promoted_args_ = True
-    args, varargs, varkw, defaults = inspect.getargspec(func)
-    args = ["v%s" % (i, ) for i in range(len(args))]
-    assert varargs is None and varkw is None
-    assert not defaults
-    argstring = ", ".join(args)
-    code = ["def f(%s):\n" % (argstring, )]
-    for arg in args:
-        code.append("    %s = hint(%s, promote=True)\n" % (arg, arg))
-    code.append("    return func(%s)\n" % (argstring, ))
-    d = {"func": func, "hint": hint}
-    exec py.code.Source("\n".join(code)).compile() in d
-    result = d["f"]
-    result.func_name = func.func_name + "_promote"
-    return result
+def purefunction_promote(promote_args='all'):
+    def decorator(func):
+        import inspect
+        purefunction(func)
+        func._pure_function_with_all_promoted_args_ = True
+        args, varargs, varkw, defaults = inspect.getargspec(func)
+        args = ["v%s" % (i, ) for i in range(len(args))]
+        assert varargs is None and varkw is None
+        assert not defaults
+        argstring = ", ".join(args)
+        code = ["def f(%s):\n" % (argstring, )]
+        if promote_args != 'all':
+            args = [('v%d' % int(i)) for i in promote_args.split(",")]
+        for arg in args:
+            code.append("    %s = hint(%s, promote=True)\n" % (arg, arg))
+        code.append("    return func(%s)\n" % (argstring, ))
+        d = {"func": func, "hint": hint}
+        exec py.code.Source("\n".join(code)).compile() in d
+        result = d["f"]
+        result.func_name = func.func_name + "_promote"
+        return result
+    return decorator
 
 class Entry(ExtRegistryEntry):
     _about_ = hint

Modified: pypy/branch/string-promote/pypy/rlib/test/test_jit.py
==============================================================================
--- pypy/branch/string-promote/pypy/rlib/test/test_jit.py	(original)
+++ pypy/branch/string-promote/pypy/rlib/test/test_jit.py	Fri Feb 19 20:40:18 2010
@@ -25,7 +25,7 @@
         assert res == 5
 
     def test_purefunction_promote(self):
-        @purefunction_promote
+        @purefunction_promote()
         def g(func):
             return func + 1
         def f(x):
@@ -33,6 +33,30 @@
         res = self.interpret(f, [2])
         assert res == 5
 
+    def test_purefunction_promote_args(self):
+        @purefunction_promote(promote_args='0')
+        def g(func, x):
+            return func + 1
+        def f(x):
+            return g(x * 2, x)
+        
+        import dis
+        from StringIO import StringIO
+        import sys
+        
+        s = StringIO()
+        sys.stdout = s
+        dis.dis(g)
+        sys.stdout = sys.__stdout__
+        x = s.getvalue().find('CALL_FUNCTION')
+        assert x != -1
+        x = s.getvalue().find('CALL_FUNCTION', x)
+        assert x != -1
+        x = s.getvalue().find('CALL_FUNCTION', x)
+        assert x != -1
+        res = self.interpret(f, [2])
+        assert res == 5
+
     def test_annotate_hooks(self):
         
         def can_inline(m): pass



More information about the Pypy-commit mailing list