[pypy-commit] pypy default: Port @with_unicode_literals from the py3k branch

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: 
Changeset: r60576:ef2d35b92279
Date: 2013-01-27 20:50 +0100
http://bitbucket.org/pypy/pypy/changeset/ef2d35b92279/

Log:	Port @with_unicode_literals from the py3k branch

diff --git a/rpython/tool/sourcetools.py b/rpython/tool/sourcetools.py
--- a/rpython/tool/sourcetools.py
+++ b/rpython/tool/sourcetools.py
@@ -6,6 +6,7 @@
 # XXX We should try to generalize and single out one approach to dynamic
 # XXX code compilation.
 
+import types
 import sys, os, inspect, new
 import py
 
@@ -295,3 +296,40 @@
     result.func_defaults = f.func_defaults
     result.func_dict.update(f.func_dict)
     return result
+
+
+def _convert_const_maybe(x, encoding):
+    if isinstance(x, str):
+        return x.decode(encoding)
+    elif isinstance(x, tuple):
+        items = [_convert_const_maybe(item, encoding) for item in x]
+        return tuple(items)
+    return x
+    
+def with_unicode_literals(fn=None, **kwds):
+    """Decorator that replace all string literals with unicode literals.
+    Similar to 'from __future__ import string literals' at function level.
+    Useful to limit changes in the py3k branch.
+    """
+    encoding = kwds.pop('encoding', 'ascii')
+    if kwds:
+        raise TypeError("Unexpected keyword argument(s): %s" % ', '.join(kwds.keys()))
+    def decorator(fn):
+        co = fn.func_code
+        new_consts = []
+        for const in co.co_consts:
+            new_consts.append(_convert_const_maybe(const, encoding))
+        new_consts = tuple(new_consts)
+        new_code = types.CodeType(co.co_argcount, co.co_nlocals, co.co_stacksize,
+                                  co.co_flags, co.co_code, new_consts, co.co_names,
+                                  co.co_varnames, co.co_filename, co.co_name,
+                                  co.co_firstlineno, co.co_lnotab)
+        fn.func_code = new_code
+        return fn
+    #
+    # support the usage of @with_unicode_literals instead of @with_unicode_literals()
+    if fn is not None:
+        assert type(fn) is types.FunctionType
+        return decorator(fn)
+    else:
+        return decorator
diff --git a/rpython/tool/test/test_sourcetools.py b/rpython/tool/test/test_sourcetools.py
--- a/rpython/tool/test/test_sourcetools.py
+++ b/rpython/tool/test/test_sourcetools.py
@@ -1,4 +1,7 @@
-from rpython.tool.sourcetools import func_with_new_name, func_renamer, rpython_wrapper
+# -*- encoding: utf-8 -*-
+import py
+from rpython.tool.sourcetools import (
+    func_with_new_name, func_renamer, rpython_wrapper, with_unicode_literals)
 
 def test_rename():
     def f(x, y=5):
@@ -56,3 +59,28 @@
         ]
 
         
+def test_with_unicode_literals():
+    @with_unicode_literals()
+    def foo():
+        return 'hello'
+    assert type(foo()) is unicode
+    #
+    @with_unicode_literals
+    def foo():
+        return 'hello'
+    assert type(foo()) is unicode
+    #
+    def foo():
+        return 'hello àèì'
+    py.test.raises(UnicodeDecodeError, "with_unicode_literals(foo)")
+    #
+    @with_unicode_literals(encoding='utf-8')
+    def foo():
+        return 'hello àèì'
+    assert foo() == u'hello àèì'
+    #
+    @with_unicode_literals
+    def foo():
+        return ('a', 'b')
+    assert type(foo()[0]) is unicode
+