pkg: comply with design of other dvc commands

Signed-off-by: Ruslan Kuprieiev <[email protected]>

dvc/command/pkg.py

@@ -5,28 +5,26 @@ import logging
 
 from dvc.exceptions import DvcException
 from dvc.command.base import CmdBase, fix_subparsers, append_doc_link
-from dvc.repo.pkg import PackageParams
 
 
 logger = logging.getLogger(__name__)
 
 
-class CmdPkg(CmdBase):
-    def run(self, unlock=False):
+class CmdPkgInstall(CmdBase):
+    def run(self):
         try:
-            pkg_params = PackageParams(
+            self.repo.pkg.install(
                 self.args.address,
                 self.args.target_dir,
                 self.args.select,
                 self.args.file,
             )
-            return self.repo.install_pkg(pkg_params)
+            return 0
         except DvcException:
             logger.exception(
                 "failed to install package '{}'".format(self.args.address)
             )
             return 1
-        pass
 
 
 def add_parser(subparsers, parent_parser):
@@ -88,4 +86,4 @@ def add_parser(subparsers, parent_parser):
         "By default the file has 'mod_' prefix and imported package name "
         "followed by .dvc",
     )
-    pkg_install_parser.set_defaults(func=CmdPkg)
+    pkg_install_parser.set_defaults(func=CmdPkgInstall)

dvc/repo/__init__.py

@@ -35,7 +35,6 @@ class Repo(object):
     from dvc.repo.status import status
     from dvc.repo.gc import gc
     from dvc.repo.commit import commit
-    from dvc.repo.pkg import install_pkg
     from dvc.repo.diff import diff
     from dvc.repo.brancher import brancher
 
@@ -50,6 +49,7 @@ class Repo(object):
         from dvc.repo.metrics import Metrics
         from dvc.scm.tree import WorkingTree
         from dvc.repo.tag import Tag
+        from dvc.repo.pkg import Pkg
 
         root_dir = self.find_root(root_dir)
 
@@ -80,6 +80,7 @@ class Repo(object):
 
         self.metrics = Metrics(self)
         self.tag = Tag(self)
+        self.pkg = Pkg(self)
 
         self._ignore()
 

dvc/repo/pkg/__init__.py

@@ -0,0 +1,8 @@
+class Pkg(object):
+    def __init__(self, repo):
+        self.repo = repo
+
+    def install(self, *args, **kwargs):
+        from dvc.repo.pkg.install import install
+
+        return install(self.repo, *args, **kwargs)

dvc/repo/pkg/install.py

@@ -12,20 +12,6 @@ from dvc.scm.git.temp_git_repo import TempGitRepo
 logger = logging.getLogger(__name__)
 
 
-class PackageParams(object):
-    def __init__(self, address, target_dir, select=[], file=None):
-        self.address = address
-        self.target_dir = target_dir
-        self.select = select
-        self.file = file
-
-    @property
-    def all_addresses(self):
-        if self.address:
-            return [self.address]
-        return PackageManager.read_packages()
-
-
 class PackageManager(object):
     PACKAGE_FILE = "package.yaml"
 
@@ -49,7 +35,9 @@ class PackageManager(object):
 class Package(object):
     MODULES_DIR = "dvc_mod"
 
-    def install_or_update(self, parent_repo, pkg_param):
+    def install_or_update(
+        self, parent_repo, address, target_dir, select=[], fname=None
+    ):
         raise NotImplementedError(
             "A method of abstract Package class was called"
         )
@@ -61,7 +49,9 @@ class Package(object):
 class GitPackage(Package):
     DEF_DVC_FILE_PREFIX = "mod_"
 
-    def install_or_update(self, parent_repo, pkg_params):
+    def install_or_update(
+        self, parent_repo, address, target_dir, select=[], fname=None
+    ):
         from git.cmd import Git
 
         if not self.is_in_root():
@@ -74,24 +64,19 @@ class GitPackage(Package):
             os.makedirs(self.MODULES_DIR)
             parent_repo.scm.ignore(os.path.abspath(self.MODULES_DIR))
 
-        module_name = (
-            Git.polish_url(pkg_params.address).strip("/").split("/")[-1]
-        )
+        module_name = Git.polish_url(address).strip("/").split("/")[-1]
         if not module_name:
             raise DvcException(
                 "Package address error: unable to extract package name"
             )
 
         with TempGitRepo(
-            pkg_params.address, module_name, Package.MODULES_DIR
+            address, module_name, Package.MODULES_DIR
         ) as tmp_repo:
             outputs_to_copy = tmp_repo.outs
-            if pkg_params.select:
+            if select:
                 outputs_to_copy = list(
-                    filter(
-                        lambda out: out.dvc_path in pkg_params.select,
-                        outputs_to_copy,
-                    )
+                    filter(lambda out: out.dvc_path in select, outputs_to_copy)
                 )
 
             fetched_stage_files = set(
@@ -102,22 +87,17 @@ class GitPackage(Package):
             module_dir = self.create_module_dir(module_name)
             tmp_repo.persist_to(module_dir, parent_repo)
 
-            dvc_file = self.get_dvc_file_name(
-                pkg_params.file, pkg_params.target_dir, module_name
-            )
+            dvc_file = self.get_dvc_file_name(fname, target_dir, module_name)
             try:
                 self.persist_stage_and_scm_state(
-                    parent_repo,
-                    outputs_to_copy,
-                    pkg_params.target_dir,
-                    dvc_file,
+                    parent_repo, outputs_to_copy, target_dir, dvc_file
                 )
             except Exception as ex:
                 raise DvcException(
                     "Package '{}' was installed "
                     "but stage file '{}' "
                     "was not created properly: {}".format(
-                        pkg_params.address, dvc_file, ex
+                        address, dvc_file, ex
                     )
                 )
 
@@ -161,7 +141,7 @@ class GitPackage(Package):
         return dvc_file_path
 
 
-def install_pkg(self, pkg_params):
+def install(self, address, target_dir, select=[], fname=None):
     """
     Install package.
 
@@ -175,31 +155,19 @@ def install_pkg(self, pkg_params):
           Result: tag_classifier package in dvc_mod/ directory
     """
 
-    if not os.path.isdir(pkg_params.target_dir):
-        logger.error(
-            "Unable to install package: "
-            "target directory '{}' does not exist".format(
-                pkg_params.target_dir
-            )
+    if not os.path.isdir(target_dir):
+        raise DvcException(
+            "target directory '{}' does not exist".format(target_dir)
         )
-        return 1
 
     curr_dir = os.path.realpath(os.curdir)
-    if not os.path.realpath(pkg_params.target_dir).startswith(curr_dir):
-        logger.error(
-            "Unable to install package: the current dir should be"
-            " a subdirectory of the target dir {}".format(
-                pkg_params.target_dir
-            )
+    if not os.path.realpath(target_dir).startswith(curr_dir):
+        raise DvcException(
+            "the current directory should be a subdirectory of the target "
+            "dir '{}'".format(target_dir)
         )
-        return 1
 
-    for addr in pkg_params.all_addresses:
-        try:
-            mgr = PackageManager.get_package(addr)
-            mgr.install_or_update(self, pkg_params)
-        except Exception as ex:
-            logger.exception("Unable to install package: ".format(ex))
-            return 1
-
-    return 0
+    addresses = [address] if address else PackageManager.read_packages()
+    for addr in addresses:
+        mgr = PackageManager.get_package(addr)
+        mgr.install_or_update(self, address, target_dir, select, fname)

tests/func/test_pkg.py

@@ -1,6 +1,6 @@
 import os
 
-from dvc.repo.pkg import Package, PackageParams
+from dvc.repo.pkg.install import Package
 from tests.basic_env import TestDvcPkg
 
 
@@ -12,9 +12,7 @@ class TestPkgBasic(TestDvcPkg):
 
         self.to_dir = os.path.join(self._root_dir, self.NEW_PKG_DIR)
         os.mkdir(self.to_dir)
-        self.dvc.install_pkg(
-            PackageParams(self.pkg_fixture._root_dir, self.to_dir)
-        )
+        self.dvc.pkg.install(self.pkg_fixture._root_dir, self.to_dir)
 
     def test_installed_dirs_and_files(self):
         self.assertTrue(os.path.isdir(self.pkg_dir))