From 21a2504bcb0ebf490827430c1581d96374bdea9f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?=
 =?UTF-8?q?=E0=A4=A4=29?= <suagatchhetri@outlook.com>
Date: Wed, 7 Jun 2023 18:44:17 +0545
Subject: [PATCH] dvc.utils.pkg: move PKG to top-level dvc

PKG has to be specified in dvc/_build.py now, and can be
imported using:

    from dvc import PKG

This is done to avoid importing from dvc/utils pkg which can
be expensive.
Also it's not an utility, so it is better that it remains at
the top-level namespace.
Note that we may have to fix the path for a few packages that we build,
like conda/snap/chocolatey.
---
 dvc/__init__.py             | 4 ++++
 dvc/_build.py               | 1 +
 dvc/cli/__init__.py         | 9 ++++-----
 dvc/info.py                 | 3 +--
 dvc/updater.py              | 3 +--
 dvc/utils/pkg.py            | 6 ------
 scripts/build.py            | 2 +-
 scripts/build_package.sh    | 2 +-
 tests/unit/cli/test_main.py | 2 +-
 9 files changed, 14 insertions(+), 18 deletions(-)
 create mode 100644 dvc/_build.py
 delete mode 100644 dvc/utils/pkg.py

diff --git a/dvc/__init__.py b/dvc/__init__.py
index 11b89b84d..22b0bcb9c 100644
--- a/dvc/__init__.py
+++ b/dvc/__init__.py
@@ -3,7 +3,11 @@ DVC
 ----
 Make your data science projects reproducible and shareable.
 """
+from typing import Optional
+
 import dvc.logger
+from dvc import _build
 from dvc.version import __version__, version_tuple  # noqa: F401
 
+PKG: "Optional[str]" = _build.PKG
 dvc.logger.setup()
diff --git a/dvc/_build.py b/dvc/_build.py
new file mode 100644
index 000000000..b0347a304
--- /dev/null
+++ b/dvc/_build.py
@@ -0,0 +1 @@
+PKG = None
diff --git a/dvc/cli/__init__.py b/dvc/cli/__init__.py
index 274b564fd..7114d9921 100644
--- a/dvc/cli/__init__.py
+++ b/dvc/cli/__init__.py
@@ -75,7 +75,7 @@ def _log_exceptions(exc: Exception) -> Optional[int]:
     from dvc.fs import AuthError, ConfigError, RemoteMissingDepsError
 
     if isinstance(exc, RemoteMissingDepsError):
-        from dvc.utils.pkg import PKG
+        from dvc import PKG
 
         proto = exc.protocol
         by_pkg = {
@@ -83,9 +83,9 @@ def _log_exceptions(exc: Exception) -> Optional[int]:
             "conda": f"conda install -c conda-forge dvc-{proto}",
         }
 
-        cmd = by_pkg.get(PKG)
-        if cmd:
+        if PKG in by_pkg:
             link = format_link("https://dvc.org/doc/install")
+            cmd = by_pkg.get(PKG)
             hint = (
                 "To install dvc with those dependencies, run:\n"
                 "\n"
@@ -190,8 +190,7 @@ def main(argv=None):  # noqa: C901, PLR0912, PLR0915
         if level and level <= logging.DEBUG:
             from platform import platform, python_implementation, python_version
 
-            from dvc import __version__
-            from dvc.utils.pkg import PKG
+            from dvc import PKG, __version__
 
             pyv = " ".join([python_implementation(), python_version()])
             pkg = f" ({PKG})" if PKG else ""
diff --git a/dvc/info.py b/dvc/info.py
index 4ee44c9f0..9a8f5332f 100644
--- a/dvc/info.py
+++ b/dvc/info.py
@@ -6,13 +6,12 @@ import platform
 
 import psutil
 
-from dvc import __version__
+from dvc import PKG, __version__
 from dvc.exceptions import NotDvcRepoError
 from dvc.fs import Schemes, generic, get_fs_cls, get_fs_config, registry
 from dvc.repo import Repo
 from dvc.scm import SCMError
 from dvc.utils import error_link
-from dvc.utils.pkg import PKG
 
 SUBPROJECTS = (
     "dvc_data",
diff --git a/dvc/updater.py b/dvc/updater.py
index fe59c2a2b..84e7beae2 100644
--- a/dvc/updater.py
+++ b/dvc/updater.py
@@ -6,8 +6,7 @@ from typing import TYPE_CHECKING, Optional
 
 from packaging import version
 
-from dvc import __version__
-from dvc.utils.pkg import PKG
+from dvc import PKG, __version__
 
 if TYPE_CHECKING:
     from dvc.ui import RichText
diff --git a/dvc/utils/pkg.py b/dvc/utils/pkg.py
deleted file mode 100644
index 3f507a568..000000000
--- a/dvc/utils/pkg.py
+++ /dev/null
@@ -1,6 +0,0 @@
-try:
-    # file is created during dvc build
-    # pylint:disable=unused-import
-    from .build import PKG  # type: ignore[import]
-except ImportError:
-    PKG = None  # type: ignore[assignment]
diff --git a/scripts/build.py b/scripts/build.py
index 2bab7372f..9c74b2f47 100644
--- a/scripts/build.py
+++ b/scripts/build.py
@@ -20,7 +20,7 @@ parser.add_argument("--apple-id-username")
 parser.add_argument("--apple-id-password")
 args = parser.parse_args()
 
-(dvc / "utils" / "build.py").write_text(f'PKG = "{args.pkg}"')
+(dvc / "_build.py").write_text(f'PKG = "{args.pkg}"')
 
 if not (dvc / "_dvc_version.py").exists():
     raise Exception("no version info found")
diff --git a/scripts/build_package.sh b/scripts/build_package.sh
index cdc6b9e35..ff1c8171f 100755
--- a/scripts/build_package.sh
+++ b/scripts/build_package.sh
@@ -8,7 +8,7 @@ if [ ! -d "dvc" ]; then
   exit 1
 fi
 
-echo 'PKG = "pip"' >dvc/utils/build.py
+echo 'PKG = "pip"' >dvc/_build.py
 
 python -m pip install build twine
 python -m build
diff --git a/tests/unit/cli/test_main.py b/tests/unit/cli/test_main.py
index 076e7e41f..7ff57bfa3 100644
--- a/tests/unit/cli/test_main.py
+++ b/tests/unit/cli/test_main.py
@@ -39,7 +39,7 @@ def test_state_pickle_errors_are_correctly_raised(tmp_dir, caplog, mocker):
 )
 def test_remote_missing_deps_are_correctly_reported(tmp_dir, caplog, mocker, pkg, msg):
     error = RemoteMissingDepsError(FileSystem(), "proto", "proto://", ["deps"])
-    mocker.patch("dvc.utils.pkg.PKG", pkg)
+    mocker.patch("dvc.PKG", pkg)
     mocker.patch(
         "dvc.cli.parse_args",
         return_value=Namespace(