Improve mediagoblin._compat.py2_unicode.

- Encode obj.__repr__() to bytestring if its type is unicode in Python 2.
- Add internal encode_to_utf8() decorator.
- Do not raise an exception if a class does not have an __str__() method,
  just warn.

1 files changed, 22 insertions, 10 deletions

diff --git a/mediagoblin/_compat.py b/mediagoblin/_compat.py
index 5a3fac53..9164d5fc 100644
--- a/mediagoblin/_compat.py
+++ b/mediagoblin/_compat.py
@@ -1,20 +1,32 @@
-from six import PY3
+import functools
+import warnings
 
-if PY3:
+import six
+
+if six.PY3:
     from email.mime.text import MIMEText
 else:
     from email.MIMEText import MIMEText
 
 
-# taken from
-# https://github.com/django/django/blob/master/django/utils/encoding.py
+def encode_to_utf8(method):
+    def wrapper(self):
+        if six.PY2 and isinstance(method(self), six.text_type):
+            return method(self).encode('utf-8')
+        return method(self)
+    functools.update_wrapper(wrapper, method, ['__name__', '__doc__'])
+    return wrapper
+
+
+# based on django.utils.encoding.python_2_unicode_compatible
 def py2_unicode(klass):
-    # TODO: Add support for __repr__
-    if not PY3:
+    if six.PY2:
         if '__str__' not in klass.__dict__:
-            raise ValueError("@py2_unicode cannot be applied "
-                             "to %s because it doesn't define __str__()." %
-                             klass.__name__)
+            warnings.warn("@py2_unicode 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')
+        klass.__str__ = encode_to_utf8(klass.__unicode__)
+        if '__repr__' in klass.__dict__:
+            klass.__repr__ = encode_to_utf8(klass.__repr__)
     return klass