diff options
| -rw-r--r-- | six.py | 19 | ||||
| -rw-r--r-- | test_six.py | 21 |
2 files changed, 40 insertions, 0 deletions
@@ -780,6 +780,25 @@ def add_metaclass(metaclass): return metaclass(cls.__name__, cls.__bases__, orig_vars) return wrapper + +def python_2_unicode_compatible(klass): + """ + A decorator that defines __unicode__ and __str__ methods under Python 2. + Under Python 3 it does nothing. + + To support Python 2 and 3 with a single code base, define a __str__ method + returning text and apply this decorator to the class. + """ + if PY2: + if '__str__' not in klass.__dict__: + raise ValueError("@python_2_unicode_compatible cannot be applied " + "to %s because it doesn't define __str__()." % + klass.__name__) + klass.__unicode__ = klass.__str__ + klass.__str__ = lambda self: self.__unicode__().encode('utf-8') + return klass + + # Complete the moves implementation. # This code is at the end of this module to speed up module loading. # Turn this module into a package. diff --git a/test_six.py b/test_six.py index 2a9ca74..4e95394 100644 --- a/test_six.py +++ b/test_six.py @@ -822,3 +822,24 @@ def test_assertRaisesRegex(): raise AssertionError('Bar') TestAssertRaisesRegex('test').test() + + +def test_python_2_unicode_compatible(): + @six.python_2_unicode_compatible + class MyTest(object): + def __str__(self): + return six.u('hello') + + def __bytes__(self): + return six.b('hello') + + my_test = MyTest() + + if six.PY2: + assert str(my_test) == six.b("hello") + assert unicode(my_test) == six.u("hello") + elif six.PY3: + assert bytes(my_test) == six.b("hello") + assert str(my_test) == six.u("hello") + + assert getattr(six.moves.builtins, 'bytes', str)(my_test) == six.b("hello") |