Collapse ProfilingCLI.implicit_command() into command()

profiling/__main__.py

@@ -48,41 +48,44 @@ class ProfilingCLI(click.Group):
 
     def __init__(self, *args, **kwargs):
         super(ProfilingCLI, self).__init__(*args, **kwargs)
-        self.command_aliases = {}
         self.implicit_command_name = None
+        self.command_name_aliases = {}
 
     def command(self, *args, **kwargs):
         """Usage::
 
-           @group.command(aliases=['ci'])
+           @cli.command(implicit=True)
+           def main():
+               ...
+
+           @cli.command(aliases=['ci'])
            def commit():
                ...
 
         """
+        implicit = kwargs.pop('implicit', False)
         aliases = kwargs.pop('aliases', None)
         decorator = super(ProfilingCLI, self).command(*args, **kwargs)
-        if aliases is None:
+        if not implicit and aliases is None:
+            # customized features not used.
             return decorator
-        def aliased_decorator(f):
+        def _decorator(f):
             cmd = decorator(f)
-            for alias in aliases:
-                self.command_aliases[alias] = cmd.name
-            return cmd
-        return aliased_decorator
-
-    def implicit_command(self, *args, **kwargs):
-        """Makes a registrar to define implicit command."""
-        def decorator(f):
-            if self.implicit_command_name is not None:
-                raise RuntimeError('Implicit command already defined')
-            cmd = self.command(*args, **kwargs)(f)
-            self.implicit_command_name = cmd.name
+            if implicit:
+                if self.implicit_command_name is not None:
+                    del self.commands[cmd.name]
+                    raise RuntimeError('Implicit command already defined')
+                self.implicit_command_name = cmd.name
+            if aliases:
+                for alias in aliases:
+                    self.command_name_aliases[alias] = cmd.name
             return cmd
-        return decorator
+        return _decorator
 
     def get_command(self, ctx, cmd_name):
+        # resolve alias.
         try:
-            cmd_name = self.command_aliases[cmd_name]
+            cmd_name = self.command_name_aliases[cmd_name]
         except KeyError:
             pass
         if self.implicit_command_name is None:
@@ -498,7 +501,7 @@ def collect_usage_pieces(self, ctx):
         return pieces
 
 
-@cli.implicit_command(cls=ProfilingCommand)
+@cli.command(cls=ProfilingCommand, implicit=True)
 @profiler_arguments
 @profiler_options
 @onetime_profiler_options

test/test_cli.py

@@ -1,7 +1,9 @@
 # -*- coding: utf-8 -*-
+import click
 from click.testing import CliRunner
+import pytest
 
-from profiling.__main__ import cli
+from profiling.__main__ import Module, ProfilingCLI, cli
 
 
 cli_runner = CliRunner()
@@ -11,3 +13,39 @@ def test_profiling_command_usage():
     for cmd in ['profile', 'live-profile', 'remote-profile']:
         r = cli_runner.invoke(cli, [cmd, '--help'])
         assert 'SCRIPT [--] [ARGV]...' in r.output
+
+
+def test_module_param_type():
+    t = Module()
+    # timeit
+    filename, code, globals_ = t.convert('timeit', None, None)
+    assert filename.endswith('timeit.py')
+    assert code.co_filename.endswith('timeit.py')
+    assert globals_['__name__'] == '__main__'
+    assert globals_['__file__'].endswith('timeit.py')
+    assert globals_['__package__'] == ''
+    # profiling.__main__
+    filename, code, globals_ = t.convert('profiling', None, None)
+    assert filename.endswith('profiling/__main__.py')
+    assert code.co_filename.endswith('profiling/__main__.py')
+    assert globals_['__name__'] == '__main__'
+    assert globals_['__file__'].endswith('profiling/__main__.py')
+    assert globals_['__package__'] == 'profiling'
+
+
+def test_customized_cli():
+    def f():
+        pass
+    cli = ProfilingCLI()
+    cli.command(name='foo', aliases=['fooo', 'foooo'])(f)
+    cli.command(name='bar', implicit=True)(f)
+    with pytest.raises(RuntimeError):
+        cli.command(name='baz', implicit=True)(f)
+    print cli.commands
+    assert len(cli.commands) == 2
+    ctx = click.Context(cli)
+    assert cli.get_command(ctx, 'foo').name == 'foo'
+    assert cli.get_command(ctx, 'fooo').name == 'foo'
+    assert cli.get_command(ctx, 'foooo').name == 'foo'
+    assert cli.get_command(ctx, 'bar').name == 'bar'
+    assert cli.get_command(ctx, 'hello.txt').name == 'bar'

test/test_timers.py

@@ -50,6 +50,7 @@ def profile(profiler):
     assert stat1.deep_time < stat2.deep_time
 
 
+@pytest.mark.flaky(reruns=10)
 def test_thread_timer():
     if sys.version_info < (3, 3):
         pytest.importorskip('yappi')