Support application-wide options

Author: jirikuncar

flask_script/__init__.py

@@ -16,6 +16,14 @@
 __all__ = ["Command", "Shell", "Server", "Manager", "Group", "Option",
            "prompt", "prompt_pass", "prompt_bool", "prompt_choices"]
 
+safe_actions = (argparse._StoreAction,
+                argparse._StoreConstAction,
+                argparse._StoreTrueAction,
+                argparse._StoreFalseAction,
+                argparse._AppendAction,
+                argparse._AppendConstAction,
+                argparse._CountAction)
+
 
 class Manager(object):
     """
@@ -131,12 +139,12 @@ def create_parser(self, prog, parser=None):
         for option in self.get_options():
             parser.add_argument(*option.args, **option.kwargs)
 
-        subparsers = parser.add_subparsers()
+        subparsers = parser.add_subparsers()  # dest='subparser_'+prog)
         for name, command in self._commands.iteritems():
             description = getattr(command, 'description',
                                   'Perform command ' + name)
             p = subparsers.add_parser(name, help=description, usage=description)
-            command.create_parser(name, parser=p)
+            p = command.create_parser(name, parser=p)
 
         return parser
 
@@ -305,34 +313,43 @@ def _handle(self, app, *args, **kwargs):
             self.print_usage()
             return 1
 
-    def handle(self, prog, name, args=None):
+    def handle(self, prog, args=None):
 
         args = list(args or [])
 
-        try:
-            command = self._commands[name]
-        except KeyError:
-            raise InvalidCommand("Command %s not found" % name)
-
         app_parser = self.create_parser(prog)
-        app_namespace, remaining_args = app_parser.parse_known_args([name] + args)
+        app_namespace, remaining_args = app_parser.parse_known_args(args)
 
         kwargs = app_namespace.__dict__
         handle = kwargs['func_handle']
         del kwargs['func_handle']
-        app = self.create_app(**app_namespace.__dict__)
 
-        # get command from bounded handle function (py2.7+)
+        ## get only safe config options
+        app_config_keys = [action.dest for action in app_parser._actions
+                           if action.__class__ in safe_actions]
+
+        ## pass only safe app config keys
+        app_config = dict((k, v) for k, v in kwargs.iteritems()
+                          if k in app_config_keys)
+
+        ## remove application config keys from handle kwargs
+        kwargs = dict((k, v) for k, v in kwargs.iteritems()
+                      if k not in app_config_keys)
+
+        ## get command from bounded handle function (py2.7+)
         command = handle.__self__
         if getattr(command, 'capture_all_args', False):
             positional_args = [remaining_args]
         else:
             if len(remaining_args):
                 # raise correct exception
                 # FIXME maybe change capture_all_args flag
-                app_parser.parse_args([name] + args)
+                app_parser.parse_args(args)
                 # sys.exit(2)
+                pass
             positional_args = []
+
+        app = self.create_app(**app_config)
         return handle(app, *positional_args, **kwargs)
 
     def run(self, commands=None, default_command=None):
@@ -351,21 +368,14 @@ def run(self, commands=None, default_command=None):
         if commands:
             self._commands.update(commands)
 
-        try:
-            try:
-                command = sys.argv[1]
-            except IndexError:
-                command = default_command
-
-            if command is None:
-                raise InvalidCommand("Please provide a command:")
-
-            result = self.handle(sys.argv[0], command, sys.argv[2:])
+        if default_command is not None and len(sys.argv) == 1:
+            sys.argv.append(default_command)
 
-            sys.exit(result or 0)
-
-        except InvalidCommand, e:
-            print e
-            self.print_usage()
+        try:
+            result = self.handle(sys.argv[0], sys.argv[1:])
+        except SystemExit as e:
+            result = e.code
+        except Exception as e:
+            raise e
 
-        sys.exit(1)
+        sys.exit(result or 0)

tests.py

@@ -24,9 +24,12 @@ def run(command_line, manager_run, capture_stderr=False):
     except SystemExit as e:
         exit_code = e.code
     finally:
+        out = sys.stdout.getvalue()
+        if capture_stderr:
+            out += sys.stderr.getvalue()
         sys.stderr = sys_stderr_orig
 
-    return sys.stdout.getvalue(), exit_code
+    return out, exit_code
 
 
 class SimpleCommand(Command):
@@ -257,23 +260,18 @@ def hello_again(name, url=None):
         stdout, code = run('manage.py hello_again --name=joe --url=reddit.com', lambda: manager.run())
         assert 'hello joe from reddit.com' in stdout
 
-    def test_global_option_provided_before_and_after_command(self):
+    def test_global_option_provided_before_command(self):
 
         manager = Manager(self.app)
         manager.add_option('-c', '--config', dest='config_name', required=False, default='Development')
         manager.add_command('simple', SimpleCommand())
 
         assert isinstance(manager._commands['simple'], SimpleCommand)
 
-        stdout, code = run('manage.py simple -c Development', lambda: manager.run())
+        stdout, code = run('manage.py -c Development simple', lambda: manager.run())
         assert code == 0
         assert 'OK' in stdout
 
-        # TODO: Allow swapping '-c Development' option and 'simple' command
-        # stdout, code = run('manage.py -c Development simple', lambda: manager.run())
-        # assert code == 0
-        # assert 'OK' in stdout
-
     def test_get_usage(self):
 
         manager = Manager(self.app)
@@ -300,7 +298,7 @@ def test_run_existing_command(self):
     def test_run_non_existant_command(self):
 
         manager = Manager(self.app)
-        self.assertRaises(InvalidCommand, manager.handle, 'manage.py', 'simple')
+        self.assertRaises(SystemExit, manager.handle, 'manage.py', 'simple')
 
     def test_run_existing(self):
 
@@ -324,15 +322,15 @@ def test_run_not_existing(self):
         manager = Manager(self.app)
 
         stdout, code = run('manage.py simple', lambda: manager.run())
-        assert code == 1
+        assert code == 2
         assert 'OK' not in stdout
 
     def test_run_no_name(self):
 
         manager = Manager(self.app)
 
         stdout, code = run('manage.py', lambda: manager.run())
-        assert code == 1
+        assert code == 2
 
     def test_run_good_options(self):
 
@@ -372,7 +370,7 @@ def test_init_with_flask_instance(self):
         assert callable(manager.app)
 
     def test_init_with_callable(self):
-        manager = Manager(lambda: app)
+        manager = Manager(lambda: self.app)
         assert callable(manager.app)
 
     def test_raise_index_error(self):
@@ -384,7 +382,7 @@ def error():
             raise IndexError()
 
         try:
-            self.assertRaises(IndexError, manager.run, default_command='error')
+            self.assertRaises(IndexError, run, 'manage.py error', lambda: manager.run())
         except SystemExit, e:
             assert e.code == 1
 
@@ -442,7 +440,7 @@ def test_submanager_has_options(self):
         assert code == 0
         assert 'OK' in stdout
 
-        stdout, code = run('manage.py sub_manager simple -c Development', lambda: manager.run())
+        stdout, code = run('manage.py -c Development sub_manager simple', lambda: manager.run())
         assert code == 0
         assert 'OK' in stdout
 
@@ -453,9 +451,9 @@ def test_manager_usage_with_submanager(self):
         manager = Manager(self.app)
         manager.add_command('sub_manager', sub_manager)
 
-        stdout, code = run('manage.py', lambda: manager.run())
-        assert code == 1
-        assert 'sub_manager  Example sub-manager' in stdout
+        stdout, code = run('manage.py -h', lambda: manager.run())
+        assert code == 0
+        assert 'Example sub-manager' in stdout
 
     def test_submanager_usage(self):
 
@@ -465,9 +463,14 @@ def test_submanager_usage(self):
         manager = Manager(self.app)
         manager.add_command('sub_manager', sub_manager)
 
-        stdout, code = run('manage.py sub_manager', lambda: manager.run())
-        assert code == 1
-        assert 'simple  simple command' in stdout
+        stdout, code = run('manage.py sub_manager', lambda: manager.run(),
+                           capture_stderr=True)
+        assert code == 2
+        assert 'too few arguments' in stdout
+
+        stdout, code = run('manage.py sub_manager -h', lambda: manager.run())
+        assert code == 0
+        assert 'simple command' in stdout
 
     def test_submanager_has_no_default_commands(self):