fix #465 and add IPython-style py3compat cast functions

thomasballinger · thomasballinger · commit 23321f24b582 · 2015-01-25T18:48:46.000-05:00
diff --git a/bpython/_py3compat.py b/bpython/_py3compat.py
@@ -35,10 +35,23 @@
 """
 
 import sys
+import locale
 
 py3 = (sys.version_info[0] == 3)
 
 if py3:
     from pygments.lexers import Python3Lexer as PythonLexer
 else:
     from pygments.lexers import PythonLexer
+
+
+def cast_unicode(s, encoding=locale.getpreferredencoding()):
+    if isinstance(s, bytes):
+        return s.decode(encoding)
+    return s
+
+
+def cast_bytes(s, encoding=locale.getpreferredencoding()):
+    if not isinstance(s, bytes):
+        return s.encode(encoding)
+    return s
diff --git a/bpython/curtsiesfrontend/repl.py b/bpython/curtsiesfrontend/repl.py
@@ -14,7 +14,7 @@
 import unicodedata
 
 from pygments import format
-from bpython._py3compat import PythonLexer
+from bpython._py3compat import PythonLexer, cast_unicode, cast_bytes
 from pygments.formatters import TerminalFormatter
 
 import blessings
@@ -1362,6 +1362,9 @@ def focus_on_subprocess(self, args):
             signal.signal(signal.SIGWINCH, prev_sigwinch_handler)
 
     def pager(self, text):
+        """Runs an external pager on text
+
+        text must be a bytestring, ie not yet encoded"""
         command = get_pager_command()
         with tempfile.NamedTemporaryFile() as tmp:
             tmp.write(text)
@@ -1375,12 +1378,14 @@ def show_source(self):
             self.status_bar.message(str(e))
         else:
             if self.config.highlight_show_source:
-                source = format(PythonLexer().get_tokens(source),
-                                TerminalFormatter())
+                source = cast_bytes(format(PythonLexer().get_tokens(source),
+                                           TerminalFormatter()))
+            else:
+                source = cast_bytes(source)
             self.pager(source)
 
     def help_text(self):
-        return (self.version_help_text() + '\n' + self.key_help_text()).encode('utf8')
+        return cast_bytes(self.version_help_text() + '\n' + self.key_help_text())
 
     def version_help_text(self):
         return (('bpython-curtsies version %s' % bpython.__version__) + ' ' +
diff --git a/bpython/repl.py b/bpython/repl.py
@@ -468,7 +468,7 @@ def get_args(self):
     def get_source_of_current_name(self):
         """Return the 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."""
+        source cannot be found. Returns bytestring in py2, unicode in py3."""
         obj = self.current_func
         try:
             if obj is None:
diff --git a/bpython/test/test_curtsies_repl.py b/bpython/test/test_curtsies_repl.py
@@ -206,7 +206,7 @@ def captured_output():
         sys.stdout, sys.stderr = old_out, old_err
 
 def create_repl(**kwargs):
-    config = setup_config({'editor':'true'})
+    config = setup_config({'editor': 'true'})
     repl = curtsiesrepl.Repl(config=config, **kwargs)
     os.environ['PAGER'] = 'true'
     repl.width = 50
@@ -264,6 +264,27 @@ def test_variable_is_cleared(self):
         self.repl.undo()
         self.assertNotIn('b', self.repl.interp.locals)
 
+class TestCurtsiesPagerText(unittest.TestCase):
+
+    def setUp(self):
+        self.repl = create_repl()
+        self.repl.pager = self.assert_pager_gets_bytes
+
+    def assert_pager_gets_bytes(self, text):
+        self.assertIsInstance(text, type(b''))
+
+    def test_help(self):
+        self.repl.pager(self.repl.help_text())
+
+    def test_show_source_not_formatted(self):
+        self.repl.config.highlight_show_source = False
+        self.repl.get_source_of_current_name = lambda: 'source code å∂ßƒåß∂ƒ'
+        self.repl.show_source()
+
+    def test_show_source_formatted(self):
+        self.repl.config.highlight_show_source = True
+        self.repl.get_source_of_current_name = lambda: 'source code å∂ßƒåß∂ƒ'
+        self.repl.show_source()
 
 if __name__ == '__main__':
     unittest.main()
