Andy Mikhailenko avatar Andy Mikhailenko committed 8a83c21

Fix #15: command output was broken in Python 3.

Comments (0)

Files changed (2)

 import sys
 from types import GeneratorType
 
-from argh.six import b, u, string_types, text_type, BytesIO, PY3
+from argh.six import (b, u, string_types, binary_type, text_type,
+                      BytesIO, StringIO, PY3)
 from argh.exceptions import CommandError
 from argh.utils import get_subparsers
 from argh.completion import autocomplete
     if output_file is None:
         # user wants a string; we create an internal temporary file-like object
         # and will return its contents as a string
-        f = BytesIO()
+        f = StringIO() if PY3 else BytesIO()
     else:
         # normally this is stdout; can be any file
         f = output_file
         # print the line as soon as it is generated to ensure that it is
         # displayed to the user before anything else happens, e.g.
         # raw_input() is called
-        output = _encode(line, f, encoding)
-        output = '' if output is None else output
+
+        output = _encode(line, f, encoding) or ''
         f.write(output)
         if not raw_output:
             # in most cases user wants on message per line
-            f.write(b('\n'))
+            f.write('\n')
 
     if output_file is None:
         # user wanted a string; return contents of our temporary file-like obj
 def _encode(line, output_file, encoding=None):
     """Converts given string to given encoding. If no encoding is specified, it
     is determined from terminal settings or, if none, from system settings.
+
+    .. note:: Compatibility
+
+       :Python 2.x:
+           `sys.stdout` is a file-like object that accepts `str` (bytes)
+           and breaks when `unicode` is passed to `sys.stdout.write()`.
+       :Python 3.x:
+           `sys.stdout` is a `_io.TextIOWrapper` instance that accepts `str`
+           (unicode) and breaks on `bytes`.
+
+       In Python 2.x arbitrary types are coerced to `unicode` and then to `str`.
+
+       In Python 3.x all types are coerced to `str` with the exception
+       for `bytes` which is **not allowed** to avoid confusion.
+
     """
-    # Convert string to Unicode
     if not isinstance(line, text_type):
+        if PY3 and isinstance(line, binary_type):
+            # in Python 3.x we require Unicode, period.
+            raise TypeError('Binary comand output is not supported '
+                            'in Python 3.x')
+
+        # in Python 2.x we accept bytes and convert them to Unicode.
         try:
             line = text_type(line)
         except UnicodeDecodeError:
             line = b(line).decode('utf-8')
 
+    if PY3:
+        return line
+
     # Choose output encoding
     if not encoding:
         # choose between terminal's and system's preferred encodings
         if output_file.isatty():
             encoding = getattr(output_file, 'encoding', None)
+
         encoding = encoding or locale.getpreferredencoding()
 
     # Convert string from Unicode to the output encoding
             False: ('y','N'),
         }
         y, n = defaults[default]
-        prompt = u('{action}? ({y}/{n})').format(**locals()).encode('utf-8')
+        prompt = u('{action}? ({y}/{n})').format(**locals())
+        if not PY3:
+            prompt = prompt.encode('utf-8')
         choice = None
         try:
             if default is None:
 
 import sys
 from argh.six import (
-    BytesIO, u, b, string_types, text_type, binary_type, iteritems
+    PY3, BytesIO, StringIO, u, string_types, text_type, binary_type,
+    iteritems
 )
 import unittest2 as unittest
 import argparse
 from argh import completion
 
 
+def make_IO():
+    # NOTE: this is according to sys.stdout
+    if PY3:
+        return StringIO()
+    else:
+        return BytesIO()
+
 class DebugArghParser(ArghParser):
     "(does not print stuff to stderr on exit)"
 
         else:
             args = command_string
 
-        io = BytesIO()
+        io = make_IO()
+
         if 'output_file' not in kwargs:
             kwargs['output_file'] = io
 
     def test_argv(self):
         _argv = sys.argv
         sys.argv = sys.argv[:1] + ['echo', 'hi there']
-        self.assert_cmd_returns(None, b('you said hi there\n'))
+        self.assert_cmd_returns(None, 'you said hi there\n')
         sys.argv = _argv
 
     def test_no_command(self):
 
     def test_echo(self):
         "A simple command is resolved to a function."
-        self.assert_cmd_returns('echo foo', b('you said foo\n'))
+        self.assert_cmd_returns('echo foo', 'you said foo\n')
 
     def test_bool_action(self):
         "Action `store_true`/`store_false` is inferred from default value."
-        self.assert_cmd_returns('echo --twice foo', b('you said fooyou said foo\n'))
+        self.assert_cmd_returns('echo --twice foo', 'you said fooyou said foo\n')
 
     def test_plain_signature(self):
         "Arguments can be passed to the function without a Namespace instance."
-        self.assert_cmd_returns('plain-echo bar', b('you said bar\n'))
+        self.assert_cmd_returns('plain-echo bar', 'you said bar\n')
 
     def test_bare_namespace(self):
         "A command can be resolved to a function, not a namespace."
 
     def test_namespaced_function(self):
         "A subcommand is resolved to a function."
-        self.assert_cmd_returns('greet hello', b('Hello world!\n'))
-        self.assert_cmd_returns('greet hello --name=John', b('Hello John!\n'))
+        self.assert_cmd_returns('greet hello', 'Hello world!\n')
+        self.assert_cmd_returns('greet hello --name=John', 'Hello John!\n')
         self.assert_cmd_fails('greet hello John', 'unrecognized arguments')
         self.assert_cmd_fails('greet howdy --name=John', 'too few arguments')
-        self.assert_cmd_returns('greet howdy John', b('Howdy John?\n'))
+        self.assert_cmd_returns('greet howdy John', 'Howdy John?\n')
 
     def test_alias(self):
-        self.assert_cmd_returns('aliased', b('ok\n'))
+        self.assert_cmd_returns('aliased', 'ok\n')
 
     def test_help_alias(self):
         self.assert_cmd_doesnt_fail('--help')
         """Positional arguments are resolved in the order in which the @arg
         decorators are defined.
         """
-        self.assert_cmd_returns('foo-bar foo bar', b('foo\nbar\n'))
+        self.assert_cmd_returns('foo-bar foo bar', 'foo\nbar\n')
 
     def test_raw_output(self):
         "If the raw_output flag is set, no extra whitespace is added"
-        self.assert_cmd_returns('foo-bar foo bar', b('foo\nbar\n'))
-        self.assert_cmd_returns('foo-bar foo bar', b('foobar'), raw_output=True)
+        self.assert_cmd_returns('foo-bar foo bar', 'foo\nbar\n')
+        self.assert_cmd_returns('foo-bar foo bar', 'foobar', raw_output=True)
 
     def test_output_file(self):
-        self.assert_cmd_returns('greet hello', b('Hello world!\n'))
-        self.assert_cmd_returns('greet hello', b('Hello world!\n'), output_file=None)
+        self.assert_cmd_returns('greet hello', 'Hello world!\n')
+        self.assert_cmd_returns('greet hello', 'Hello world!\n', output_file=None)
 
     def test_command_error(self):
-        self.assert_cmd_returns('whiner-plain', b('I feel depressed.\n'))
-        self.assert_cmd_returns('whiner-iterable', b('Hello...\nI feel depressed.\n'))
+        self.assert_cmd_returns('whiner-plain', 'I feel depressed.\n')
+        self.assert_cmd_returns('whiner-iterable', 'Hello...\nI feel depressed.\n')
 
     def test_custom_namespace(self):
         namespace = argparse.Namespace()
         namespace.custom_value = 'foo'
-        self.assert_cmd_returns('custom-namespace', b('foo\n'),
+        self.assert_cmd_returns('custom-namespace', 'foo\n',
                                 namespace=namespace)
 
 
     def test_command_decorator(self):
         """The @command decorator creates arguments from function signature.
         """
-        self.assert_cmd_returns('command-deco', b('Hello\n'))
-        self.assert_cmd_returns('command-deco --text=hi', b('hi\n'))
+        self.assert_cmd_returns('command-deco', 'Hello\n')
+        self.assert_cmd_returns('command-deco --text=hi', 'hi\n')
 
     def test_regression_issue12(self):
         """Issue #12: @command was broken if there were more than one argument
         to begin with same character (i.e. short option names were inferred
         incorrectly).
         """
-        self.assert_cmd_returns('command-deco-issue12', b('foo 1, fox 2\n'))
-        self.assert_cmd_returns('command-deco-issue12 --foo 3', b('foo 3, fox 2\n'))
-        self.assert_cmd_returns('command-deco-issue12 --fox 3', b('foo 1, fox 3\n'))
+        self.assert_cmd_returns('command-deco-issue12', 'foo 1, fox 2\n')
+        self.assert_cmd_returns('command-deco-issue12 --foo 3', 'foo 3, fox 2\n')
+        self.assert_cmd_returns('command-deco-issue12 --fox 3', 'foo 1, fox 3\n')
         self.assert_cmd_fails('command-deco-issue12 -f 3', 'unrecognized')
 
 
         self.assertRaisesRegexp(AssertionError, 'Do it yourself', f)
 
     def test_error_wrapped(self):
-        self.assert_cmd_returns('strict-hello-smart John', b('Do it yourself\n'))
-        self.assert_cmd_returns('strict-hello-smart world', b('Hello world\n'))
+        self.assert_cmd_returns('strict-hello-smart John', 'Do it yourself\n')
+        self.assert_cmd_returns('strict-hello-smart world', 'Hello world\n')
 
 
 class NoCommandsTestCase(BaseArghTestCase):
     "Edge case: no commands defined"
     commands = {}
     def test_no_command(self):
-        self.assert_cmd_returns('', b(self.parser.format_usage()), raw_output=True)
-        self.assert_cmd_returns('', b(self.parser.format_usage()+'\n'))
+        self.assert_cmd_returns('', self.parser.format_usage(), raw_output=True)
+        self.assert_cmd_returns('', self.parser.format_usage()+'\n')
 
 
 class DefaultCommandTestCase(BaseArghTestCase):
         self.parser.set_default_command(main)
 
     def test_default_command(self):
-        self.assert_cmd_returns('', b('1\n'))
-        self.assert_cmd_returns('--foo 2', b('2\n'))
+        self.assert_cmd_returns('', '1\n')
+        self.assert_cmd_returns('--foo 2', '2\n')
         self.assert_cmd_exits('--help')
 
     def test_prevent_conflict_with_single_command(self):
         else:
             args = command_string
 
-        io = BytesIO()
+        io = make_IO()
         if 'output_file' not in kwargs:
             kwargs['output_file'] = io
 
         def main(args):
             return args.foo
 
-        self.assert_cmd_returns(main, '', b('1\n'))
-        self.assert_cmd_returns(main, '--foo 2', b('2\n'))
+        self.assert_cmd_returns(main, '', '1\n')
+        self.assert_cmd_returns(main, '--foo 2', '2\n')
 
 
 class DispatchCommandsTestCase(BaseArghTestCase):
         else:
             args = command_string
 
-        io = BytesIO()
+        io = make_IO()
         if 'output_file' not in kwargs:
             kwargs['output_file'] = io
 
         def bar(args):
             return args.y
 
-        self.assert_cmd_returns([foo, bar], 'foo', b('1\n'))
-        self.assert_cmd_returns([foo, bar], 'foo -x 5', b('5\n'))
-        self.assert_cmd_returns([foo, bar], 'bar', b('2\n'))
+        self.assert_cmd_returns([foo, bar], 'foo', '1\n')
+        self.assert_cmd_returns([foo, bar], 'foo -x 5', '5\n')
+        self.assert_cmd_returns([foo, bar], 'bar', '2\n')
 
 
 class ConfirmTestCase(unittest.TestCase):
         argh.helpers.raw_input = raw_input_mock
 
         argh.confirm('do smth')
-        self.assertEqual(prompts[-1], b('do smth? (y/n)'))
+        self.assertEqual(prompts[-1], 'do smth? (y/n)')
 
         argh.confirm('do smth', default=None)
-        self.assertEqual(prompts[-1], b('do smth? (y/n)'))
+        self.assertEqual(prompts[-1], 'do smth? (y/n)')
 
         argh.confirm('do smth', default=True)
-        self.assertEqual(prompts[-1], b('do smth? (Y/n)'))
+        self.assertEqual(prompts[-1], 'do smth? (Y/n)')
 
         argh.confirm('do smth', default=False)
-        self.assertEqual(prompts[-1], b('do smth? (y/N)'))
+        self.assertEqual(prompts[-1], 'do smth? (y/N)')
 
     def test_encoding(self):
         "Unicode and bytes are accepted as prompt message"
         def raw_input_mock(prompt):
-            assert isinstance(prompt, binary_type)
+            if not PY3:
+                assert isinstance(prompt, binary_type)
         argh.helpers.raw_input = raw_input_mock
         argh.confirm(u('привет'))
 
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.