get source in unicode in python 2 and 3

thomasballinger · thomasballinger · commit b23e1915a5fa · 2015-01-28T20:56:38.000-05:00
diff --git a/bpython/inspection.py b/bpython/inspection.py
@@ -270,3 +270,18 @@ def is_callable(obj):
         return isinstance(obj, collections.Callable)
     else:
         return callable(obj)
+
+
+def get_encoding(obj):
+    for line in inspect.getsourcelines(obj)[0][:2]:
+        m = re.search(r'coding[:=]\s*([-\w.]+)', line)
+        if m:
+            return m.group(1)
+    return 'ascii'
+
+
+def get_source_unicode(obj):
+    """Returns a decoded source of object"""
+    if py3:
+        return inspect.getsource(obj)
+    return inspect.getsource(obj).decode(get_encoding(obj))
diff --git a/bpython/repl.py b/bpython/repl.py
@@ -466,9 +466,10 @@ def get_args(self):
         return False
 
     def get_source_of_current_name(self):
-        """Return the source code of the object which is bound to the
+        """Return the unicode source code of the object which is bound to the
         current name in the current input line. Throw `SourceNotFound` if the
-        source cannot be found. Returns bytestring in py2, unicode in py3."""
+        source cannot be found."""
+
         obj = self.current_func
         try:
             if obj is None:
@@ -477,7 +478,7 @@ def get_source_of_current_name(self):
                     raise SourceNotFound(_("Nothing to get source of"))
                 if inspection.is_eval_safe_name(line):
                     obj = self.get_object(line)
-            return inspect.getsource(obj)
+            return inspection.get_source_unicode(obj)
         except (AttributeError, NameError) as e:
             msg = _("Cannot get source: %s") % (str(e), )
         except IOError as e:
diff --git a/bpython/test/test_inspection.py b/bpython/test/test_inspection.py
@@ -61,5 +61,11 @@ def spam(eggs=23, foobar="yay"):
         self.assertEqual(repr(defaults[0]), "23")
         self.assertEqual(repr(defaults[1]), "'yay'")
 
+    def test_get_encoding(self):
+        self.assertEqual(inspection.get_encoding(inspection), 'ascii')
+        from bpython.test import test_curtsies_painting
+        self.assertEqual(inspection.get_encoding(test_curtsies_painting), 'utf8')
+
+
 if __name__ == '__main__':
     unittest.main()
