diff --git a/dvc/command/base.py b/dvc/command/base.py index c70e99ad9..94e1e40d1 100644 --- a/dvc/command/base.py +++ b/dvc/command/base.py @@ -45,6 +45,10 @@ class CmdBase(ABC): updater = Updater(self.repo.tmp_dir, hardlink_lock=hardlink_lock) updater.check() + def do_run(self): + with self.repo: + return self.run() + @abstractmethod def run(self): pass @@ -55,3 +59,6 @@ class CmdBaseNoRepo(CmdBase): self.args = args os.chdir(args.cd) + + def do_run(self): + return self.run() diff --git a/dvc/command/init.py b/dvc/command/init.py index 6470554a1..a583a62cd 100644 --- a/dvc/command/init.py +++ b/dvc/command/init.py @@ -39,14 +39,14 @@ class CmdInit(CmdBaseNoRepo): from dvc.repo import Repo try: - self.repo = Repo.init( + with Repo.init( ".", no_scm=self.args.no_scm, force=self.args.force, subdir=self.args.subdir, - ) - self.config = self.repo.config - _welcome_message() + ) as repo: + self.config = repo.config + _welcome_message() except InitError: logger.exception("failed to initiate DVC") return 1 diff --git a/dvc/info.py b/dvc/info.py index 335265f5d..fdd276c29 100644 --- a/dvc/info.py +++ b/dvc/info.py @@ -31,34 +31,32 @@ def get_dvc_info(): ] try: - repo = Repo() + with Repo() as repo: + # cache_dir might not exist yet (e.g. after `dvc init`), and we + # can't auto-create it, as it might cause issues if the user + # later decides to enable shared cache mode with + # `dvc config cache.shared group`. + if os.path.exists(repo.odb.local.cache_dir): + info.append( + "Cache types: {}".format(_get_linktype_support_info(repo)) + ) + fs_type = get_fs_type(repo.odb.local.cache_dir) + info.append(f"Cache directory: {fs_type}") + else: + info.append("Cache types: " + error_link("no-dvc-cache")) - # cache_dir might not exist yet (e.g. after `dvc init`), and we - # can't auto-create it, as it might cause issues if the user - # later decides to enable shared cache mode with - # `dvc config cache.shared group`. - if os.path.exists(repo.odb.local.cache_dir): - info.append( - "Cache types: {}".format(_get_linktype_support_info(repo)) - ) - fs_type = get_fs_type(repo.odb.local.cache_dir) - info.append(f"Cache directory: {fs_type}") - else: - info.append("Cache types: " + error_link("no-dvc-cache")) - - info.append(f"Caches: {_get_caches(repo.odb)}") - - info.append(f"Remotes: {_get_remotes(repo.config)}") + info.append(f"Caches: {_get_caches(repo.odb)}") + info.append(f"Remotes: {_get_remotes(repo.config)}") + root_directory = repo.root_dir + fs_root = get_fs_type(os.path.abspath(root_directory)) + info.append(f"Workspace directory: {fs_root}") + info.append("Repo: {}".format(_get_dvc_repo_info(repo))) except NotDvcRepoError: pass except SCMError: info.append("Repo: dvc, git (broken)") - else: - root_directory = repo.root_dir - fs_root = get_fs_type(os.path.abspath(root_directory)) - info.append(f"Workspace directory: {fs_root}") - info.append("Repo: {}".format(_get_dvc_repo_info(repo))) + return "\n".join(info) diff --git a/dvc/main.py b/dvc/main.py index d0f1d3ecc..20b2326ad 100644 --- a/dvc/main.py +++ b/dvc/main.py @@ -52,7 +52,7 @@ def main(argv=None): # noqa: C901 with debugtools(args): cmd = args.func(args) - ret = cmd.run() + ret = cmd.do_run() except ConfigError: logger.exception("configuration error") ret = 251 diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index 61ed75150..961a447b2 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -209,22 +209,19 @@ class Repo: return self.url or self.root_dir @staticmethod - @contextmanager def open(url, *args, **kwargs): if url is None: url = os.getcwd() if os.path.exists(url): try: - yield Repo(url, *args, **kwargs) - return + return Repo(url, *args, **kwargs) except NotDvcRepoError: pass # fallthrough to external_repo from dvc.external_repo import external_repo - with external_repo(url, *args, **kwargs) as repo: - yield repo + return external_repo(url, *args, **kwargs) @cached_property def scm(self): @@ -307,8 +304,9 @@ class Repo: def init(root_dir=os.curdir, no_scm=False, force=False, subdir=False): from dvc.repo.init import init - init(root_dir=root_dir, no_scm=no_scm, force=force, subdir=subdir) - return Repo(root_dir) + return init( + root_dir=root_dir, no_scm=no_scm, force=force, subdir=subdir + ) def unprotect(self, target): return self.odb.local.unprotect(PathInfo(target)) @@ -508,3 +506,10 @@ class Repo: self.__dict__.pop("graph", None) self.__dict__.pop("stages", None) self.__dict__.pop("pipelines", None) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self._reset() + self.scm.close() diff --git a/tests/basic_env.py b/tests/basic_env.py index 539d28d7e..ea1b0530b 100644 --- a/tests/basic_env.py +++ b/tests/basic_env.py @@ -154,6 +154,10 @@ class TestDvcFixture(TestDirFixture): super().setUp() self.dvc = DvcRepo.init(self.root_dir, no_scm=True) + def tearDown(self): + self.dvc.close() + super().tearDown() + class TestDvcGitFixture(TestGitFixture): def setUp(self): @@ -162,7 +166,7 @@ class TestDvcGitFixture(TestGitFixture): self.dvc.scm.commit("init dvc") def tearDown(self): - self.dvc.scm.close() + self.dvc.close() super().tearDown() diff --git a/tests/dir_helpers.py b/tests/dir_helpers.py index c62ac0764..45f9e82d5 100644 --- a/tests/dir_helpers.py +++ b/tests/dir_helpers.py @@ -297,7 +297,8 @@ def scm(tmp_dir): @pytest.fixture def dvc(tmp_dir): - return tmp_dir.dvc + with tmp_dir.dvc as _dvc: + yield _dvc def git_init(path): diff --git a/tests/func/test_add.py b/tests/func/test_add.py index 6ae387e52..bda8d7f22 100644 --- a/tests/func/test_add.py +++ b/tests/func/test_add.py @@ -25,7 +25,6 @@ from dvc.hash_info import HashInfo from dvc.main import main from dvc.objects.db import ODBManager from dvc.output.base import OutputAlreadyTrackedError, OutputIsStageFileError -from dvc.repo import Repo as DvcRepo from dvc.stage import Stage from dvc.stage.exceptions import ( StageExternalOutputsError, @@ -414,7 +413,7 @@ class TestAddLocalRemoteFile(TestDvc): ret = main(["remote", "add", remote, cwd]) self.assertEqual(ret, 0) - self.dvc = DvcRepo() + self.dvc.config.load() foo = f"remote://{remote}/{self.FOO}" ret = main(["add", foo]) diff --git a/tests/func/test_checkout.py b/tests/func/test_checkout.py index e381633c9..2939dd45b 100644 --- a/tests/func/test_checkout.py +++ b/tests/func/test_checkout.py @@ -19,7 +19,6 @@ from dvc.exceptions import ( ) from dvc.fs.local import LocalFileSystem from dvc.main import main -from dvc.repo import Repo as DvcRepo from dvc.stage import Stage from dvc.stage.exceptions import StageFileDoesNotExistError from dvc.system import System @@ -92,7 +91,8 @@ class TestCheckoutCorruptedCacheDir(TestDvc): ret = main(["config", "cache.type", "copy"]) self.assertEqual(ret, 0) - self.dvc = DvcRepo(".") + self.dvc.config.load() + stages = self.dvc.add(self.DATA_DIR) self.assertEqual(len(stages), 1) self.assertEqual(len(stages[0].outs), 1) diff --git a/tests/func/test_cli.py b/tests/func/test_cli.py index d10ab1d48..079276155 100644 --- a/tests/func/test_cli.py +++ b/tests/func/test_cli.py @@ -60,31 +60,42 @@ class TestRun(TestDvc): arg2, ] ) - - self.assertIsInstance(args.func(args), CmdRun) + cmd_cls = args.func(args) + self.assertIsInstance(cmd_cls, CmdRun) self.assertEqual(args.deps, [dep1, dep2]) self.assertEqual(args.outs, [out1, out2]) self.assertEqual(args.outs_no_cache, [out_no_cache1, out_no_cache2]) self.assertEqual(args.file, fname) self.assertEqual(args.cmd, [cmd, arg1, arg2]) + cmd_cls.repo.close() + class TestPull(TestDvc): def test(self): args = parse_args(["pull"]) - self.assertIsInstance(args.func(args), CmdDataPull) + cmd = args.func(args) + self.assertIsInstance(cmd, CmdDataPull) + + cmd.repo.close() class TestPush(TestDvc): def test(self): args = parse_args(["push"]) - self.assertIsInstance(args.func(args), CmdDataPush) + cmd = args.func(args) + self.assertIsInstance(cmd, CmdDataPush) + + cmd.repo.close() class TestStatus(TestDvc): def test(self): args = parse_args(["status"]) - self.assertIsInstance(args.func(args), CmdDataStatus) + cmd = args.func(args) + self.assertIsInstance(cmd, CmdDataStatus) + + cmd.repo.close() class TestRepro(TestDvc): @@ -96,11 +107,14 @@ class TestRepro(TestDvc): ["repro", target1, target2, "-f", "--force", "-s", "--single-item"] ) - self.assertIsInstance(args.func(args), CmdRepro) + cmd = args.func(args) + self.assertIsInstance(cmd, CmdRepro) self.assertEqual(args.targets, [target1, target2]) self.assertEqual(args.force, True) self.assertEqual(args.single_item, True) + cmd.repo.close() + class TestRemove(TestDvc): def test(self): @@ -109,9 +123,12 @@ class TestRemove(TestDvc): args = parse_args(["remove", target1, target2]) - self.assertIsInstance(args.func(args), CmdRemove) + cmd = args.func(args) + self.assertIsInstance(cmd, CmdRemove) self.assertEqual(args.targets, [target1, target2]) + cmd.repo.close() + class TestAdd(TestDvc): def test(self): @@ -120,24 +137,33 @@ class TestAdd(TestDvc): args = parse_args(["add", target1, target2]) - self.assertIsInstance(args.func(args), CmdAdd) + cmd = args.func(args) + self.assertIsInstance(cmd, CmdAdd) self.assertEqual(args.targets, [target1, target2]) + cmd.repo.close() + class TestGC(TestDvc): def test(self): args = parse_args(["gc"]) - self.assertIsInstance(args.func(args), CmdGC) + cmd = args.func(args) + self.assertIsInstance(cmd, CmdGC) + + cmd.repo.close() class TestGCMultipleDvcRepos(TestDvc): def test(self): args = parse_args(["gc", "-p", "/tmp/asdf", "/tmp/xyz"]) - self.assertIsInstance(args.func(args), CmdGC) + cmd = args.func(args) + self.assertIsInstance(cmd, CmdGC) self.assertEqual(args.repos, ["/tmp/asdf", "/tmp/xyz"]) + cmd.repo.close() + class TestConfig(TestDvc): def test(self): @@ -146,7 +172,8 @@ class TestConfig(TestDvc): args = parse_args(["config", "-u", "--unset", name, value]) - self.assertIsInstance(args.func(args), CmdConfig) + cmd = args.func(args) + self.assertIsInstance(cmd, CmdConfig) self.assertEqual(args.unset, True) self.assertEqual(args.name, (False, "section", "option")) self.assertEqual(args.value, value) @@ -162,7 +189,10 @@ class TestConfig(TestDvc): class TestCheckout(TestDvc): def test(self): args = parse_args(["checkout"]) - self.assertIsInstance(args.func(args), CmdCheckout) + cmd = args.func(args) + self.assertIsInstance(cmd, CmdCheckout) + + cmd.repo.close() class TestFindRoot(TestDvc): diff --git a/tests/func/test_run_single_stage.py b/tests/func/test_run_single_stage.py index 8f03e1eb2..40afb377f 100644 --- a/tests/func/test_run_single_stage.py +++ b/tests/func/test_run_single_stage.py @@ -898,6 +898,8 @@ class TestShouldNotCheckoutUponCorruptedLocalHardlinkCache(TestDvc): super().setUp() ret = main(["config", "cache.type", "hardlink"]) self.assertEqual(ret, 0) + self.dvc.close() + self.dvc = DvcRepo(".") def test(self):