diff --git a/komandr.py b/komandr.py index 96f505b..7cb52e0 100644 --- a/komandr.py +++ b/komandr.py @@ -7,7 +7,7 @@ import inspect import argparse from functools import wraps -from itertools import izip_longest +from itertools import izip_longest, chain class prog(object): @@ -102,15 +102,27 @@ def execute(self, arg_list): :param type: list """ + if getattr(self, 'default_subcommand', None): + subcommand = arg_list[0] if arg_list else '' + action_option_strings = chain.from_iterable(i.option_strings + for i in self.parser._actions) + if not (subcommand + and subcommand in self.subparsers.choices + or subcommand in action_option_strings): + arg_list[:0] = [self.default_subcommand] + arg_map = self.parser.parse_args(arg_list).__dict__ command = arg_map.pop(self._COMMAND_FLAG) return command(**arg_map) - def __call__(self): + def __call__(self, default_subcommand=None): """Calls :py:func:``execute`` with :py:class:``sys.argv`` excluding script name which comes first. + :param default: name of subcommand called if no subcommand specified. + :param type: str """ + self.default_subcommand = default_subcommand self.execute(sys.argv[1:]) main = prog() diff --git a/tests.py b/tests.py index 59687c0..bf0c7d0 100644 --- a/tests.py +++ b/tests.py @@ -1,10 +1,14 @@ # TODO: Add much more tests import unittest import komandr +from mock import patch class TestKomandr(unittest.TestCase): + def tearDown(self): + komandr.main.default_subcommand = None + def testCommand(self): def foo(bar, baz=None): return bar, baz @@ -13,5 +17,44 @@ def foo(bar, baz=None): komandr.execute(['foo', '1', '--baz', '2'])) + def testDefaultSubcommand(self): + def foo(bar, baz=None): + return bar, baz + komandr.command(foo) + + # Without main.default_subcommand a `SystemExit` exception would raise. + with self.assertRaises(SystemExit) as ex: + komandr.execute([]) + + self.assertEqual(2, ex.exception.code) + + # With main.default_subcommand='foo', subcommand named `foo` would be + # called by default when no subcommand. + komandr.main.default_subcommand = 'foo' + self.assertEqual(('1', '2'), komandr.execute(['1', '--baz', '2'])) + + def testDefaultSubcommandWithoutOptions(self): + def foo(): + return 'bar' + komandr.command(foo) + + komandr.main.default_subcommand = 'foo' + self.assertEqual('bar', komandr.execute([])) + + def testOptionsNotTriggerDefaultSubcommand(self): + def foo(baz='bar'): + return baz + komandr.command(foo) + + komandr.main.default_subcommand = 'foo' + with patch.object(komandr.main.parser, 'print_help') as mock: + with self.assertRaises(SystemExit) as ex: + komandr.execute(['--help']) + + self.assertEqual(0, ex.exception.code) + self.assertTrue(mock.called) + + self.assertEqual('bar baz', komandr.execute(['--baz', 'bar baz'])) + if __name__ == '__main__': unittest.main()