diff --git a/dvc/cli/parser.py b/dvc/cli/parser.py index 007908c4e0..c511d11a0b 100644 --- a/dvc/cli/parser.py +++ b/dvc/cli/parser.py @@ -200,6 +200,13 @@ def get_main_parser(): type=str, ) + parser.add_argument( + "--plugins", + metavar="", + default=None, + help="Path to a plugins directory", + ) + # Sub commands subparsers = parser.add_subparsers( title="Available Commands", @@ -215,4 +222,9 @@ def get_main_parser(): for cmd in COMMANDS: cmd.add_parser(subparsers, parent_parser) + from dvc.plugins import plugin_manager + + plugin_manager.hook.register_command( # pylint: disable=no-member + parser=subparsers, parent=parent_parser + ) return parser diff --git a/dvc/env.py b/dvc/env.py index 2839772c68..16f8383a47 100644 --- a/dvc/env.py +++ b/dvc/env.py @@ -10,3 +10,4 @@ DVC_EXP_GIT_REMOTE = "DVC_EXP_GIT_REMOTE" DVC_EXP_AUTO_PUSH = "DVC_EXP_AUTO_PUSH" DVC_NO_ANALYTICS = "DVC_NO_ANALYTICS" +DVC_PLUGINS_DIR = "DVC_PLUGINS_DIR" diff --git a/dvc/hookspecs.py b/dvc/hookspecs.py new file mode 100644 index 0000000000..ed992bdc99 --- /dev/null +++ b/dvc/hookspecs.py @@ -0,0 +1,9 @@ +from pluggy import HookimplMarker, HookspecMarker + +hookspec = HookspecMarker("dvc") +hookimpl = HookimplMarker("dvc") + + +@hookspec +def register_command(parser, parent): # pylint: disable=unused-argument + pass diff --git a/dvc/plugins.py b/dvc/plugins.py new file mode 100644 index 0000000000..79d8a6938e --- /dev/null +++ b/dvc/plugins.py @@ -0,0 +1,66 @@ +import glob +import os +import sys +from typing import List + +import pluggy + +from . import hookspecs +from .env import DVC_PLUGINS_DIR + + +def import_file(name, file): + import importlib.util + + spec = importlib.util.spec_from_file_location(name, file) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +class PluginManager(pluggy.PluginManager): + def add_files_from_dir(self, directory: str) -> None: + if not os.path.isdir(directory): + return + path = os.path.join(directory, "*.py") + for file in filter(os.path.isfile, glob.glob(path)): + mod = import_file(os.path.basename(file)[:-3], file) + try: + self.register(mod) + except ValueError: + pass + + def load_from_env(self, env: str = DVC_PLUGINS_DIR) -> None: + if env in os.environ: + self.add_files_from_dir(os.environ[env]) + + def load_from_args( + self, args: List[str] = None, opt_name: str = "--plugins" + ) -> None: + args = sys.argv[1:] if args is None else args + i = 0 + n = len(args) + while i < n: + opt = args[i] + i += 1 + if isinstance(opt, str): + if opt == opt_name: + try: + parg = args[i] + except IndexError: + return + i += 1 + elif opt.startswith(f"{opt_name}="): + parg = opt[len(f"{opt_name}=") :] + else: + continue + self.add_files_from_dir(parg) + + +plugin_manager = PluginManager("dvc") +plugin_manager.add_hookspecs(hookspecs) + +if "DVC_TEST" not in os.environ: # we don't want to run this on tests + plugin_manager.load_setuptools_entrypoints("dvc") + plugin_manager.load_from_args() + plugin_manager.load_from_env() diff --git a/plugins/hello.py b/plugins/hello.py new file mode 100644 index 0000000000..3e698829df --- /dev/null +++ b/plugins/hello.py @@ -0,0 +1,21 @@ +from argparse import Namespace +from dataclasses import dataclass + +from dvc.hookspecs import hookimpl + + +@dataclass +class HelloCommand: + args: Namespace + + def do_run(self): + print("Hello, world!") + return 0 + + +@hookimpl +def register_command(parser, parent): + hello_world = parser.add_parser( + "hello", parents=[parent], help="print hello world" + ) + hello_world.set_defaults(func=HelloCommand) diff --git a/setup.cfg b/setup.cfg index 9ea3fac015..b49a013cdd 100644 --- a/setup.cfg +++ b/setup.cfg @@ -76,6 +76,7 @@ install_requires = diskcache>=5.2.1 jaraco.windows>=5.7.0; python_version < '3.8' and sys_platform == 'win32' scmrepo==0.0.7 + pluggy>=1,<2 [options.extras_require] all =