cli.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. # Copyright (c) OpenMMLab. All rights reserved.
  2. import difflib
  3. import os
  4. import os.path as osp
  5. from configparser import ConfigParser
  6. import click
  7. plugin_folder = os.path.join(os.path.dirname(__file__), 'commands')
  8. CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])
  9. DEFAULT_CFG = osp.join(osp.expanduser('~'), '.mimrc')
  10. def configure(ctx, param, filename):
  11. cfg = ConfigParser()
  12. cfg.read(filename)
  13. ctx.default_map = {}
  14. for sect in cfg.sections():
  15. command_path = sect.split('.')
  16. if command_path[0] != 'options':
  17. continue
  18. defaults = ctx.default_map
  19. for cmdname in command_path[1:]:
  20. defaults = defaults.setdefault(cmdname, {})
  21. defaults.update(cfg[sect])
  22. class MIM(click.MultiCommand):
  23. def list_commands(self, ctx):
  24. rv = []
  25. for filename in os.listdir(plugin_folder):
  26. if not filename.startswith('__') and filename.endswith('.py'):
  27. rv.append(filename[:-3])
  28. rv.sort()
  29. return rv
  30. def get_command(self, ctx, name):
  31. ns = {}
  32. fn = osp.join(plugin_folder, name + '.py')
  33. if not osp.exists(fn):
  34. matches = [
  35. x for x in self.list_commands(ctx) if x.startswith(name)
  36. ]
  37. if not matches:
  38. return None
  39. elif len(matches) == 1:
  40. return self.get_command(ctx, matches[0])
  41. with open(fn) as f:
  42. code = compile(f.read(), fn, 'exec')
  43. eval(code, ns, ns)
  44. return ns['cli']
  45. def resolve_command(self, ctx, args):
  46. # The implementation is modified from https://github.com/click-contrib/
  47. # click-didyoumean/blob/master/click_didyoumean/__init__.py#L25
  48. try:
  49. return super().resolve_command(ctx, args)
  50. except click.exceptions.UsageError as error:
  51. error_msg = str(error)
  52. original_cmd_name = click.utils.make_str(args[0])
  53. matches = difflib.get_close_matches(original_cmd_name,
  54. self.list_commands(ctx), 3,
  55. 0.1)
  56. if matches:
  57. error_msg += '\n\nDid you mean one of these?\n'
  58. error_msg += '\n'.join(matches)
  59. raise click.exceptions.UsageError(error_msg, error.ctx)
  60. @click.command(cls=MIM, context_settings=CONTEXT_SETTINGS)
  61. @click.option(
  62. '--user-conf',
  63. type=click.Path(dir_okay=False),
  64. default=DEFAULT_CFG,
  65. callback=configure,
  66. is_eager=True,
  67. expose_value=False,
  68. help='Read option defaults from the .mimrc file',
  69. show_default=True)
  70. @click.version_option()
  71. def cli():
  72. """OpenMMLab Command Line Interface.
  73. MIM provides a unified API for launching and installing OpenMMLab projects
  74. and their extensions, and managing the OpenMMLab save_models zoo.
  75. """
  76. pass