Merge pull request #2525 from efiop/2506

dvc: remove PATH prefixes appended by pyenv

dvc/utils/__init__.py

@@ -245,9 +245,11 @@ def is_binary():
     return getattr(sys, "frozen", False)
 
 
-# NOTE: Fix env variables modified by PyInstaller
-# http://pyinstaller.readthedocs.io/en/stable/runtime-information.html
 def fix_env(env=None):
+    """Fix env variables modified by PyInstaller [1] and pyenv [2].
+    [1] http://pyinstaller.readthedocs.io/en/stable/runtime-information.html
+    [2] https://github.com/pyenv/pyenv/issues/985
+    """
     if env is None:
         env = os.environ.copy()
     else:
@@ -262,6 +264,32 @@ def fix_env(env=None):
         else:
             env.pop(lp_key, None)
 
+    # Unlike PyInstaller, pyenv doesn't leave backups of original env vars
+    # when it modifies them. If we look into the shim, pyenv and pyenv-exec,
+    # we can figure out that the PATH is modified like this:
+    #
+    #     PATH=$PYENV_BIN_PATH:${bin_path}:${plugin_bin}:$PATH
+    #
+    # where
+    #
+    #     PYENV_BIN_PATH - starts with $PYENV_ROOT, see pyenv-exec source code.
+    #     bin_path - might not start with $PYENV_ROOT as it runs realpath on
+    #         it, see pyenv source code.
+    #     plugin_bin - might contain more than 1 entry, which start with
+    #         $PYENV_ROOT, see pyenv source code.
+    #
+    # So having this, we can make a rightful assumption about what parts of the
+    # PATH we need to remove in order to get the original PATH.
+    path = env.get("PATH")
+    pyenv_root = env.get("PYENV_ROOT")
+    if path and pyenv_root and path.startswith(pyenv_root):
+        # removing PYENV_BIN_PATH and bin_path
+        parts = path.split(":")[2:]
+        # removing plugin_bin from the left
+        while pyenv_root in parts[0]:
+            del parts[0]
+        env["PATH"] = ":".join(parts)
+
     return env
 
 

tests/unit/utils/test_utils.py

@@ -1,6 +1,6 @@
 import pytest
 
-from dvc.utils import to_chunks
+from dvc.utils import to_chunks, fix_env
 
 
 @pytest.mark.parametrize(
@@ -28,3 +28,30 @@ def test_to_chunks_num_chunks(num_chunks, expected_chunks):
     list_to_chunk = [1, 2, 3, 4]
     result = to_chunks(list_to_chunk, num_chunks=num_chunks)
     assert result == expected_chunks
+
+
+@pytest.mark.parametrize(
+    "path, orig",
+    [
+        (
+            "/pyenv/bin:/pyenv/libexec:/pyenv/hook:/orig/path1:/orig/path2",
+            "/orig/path1:/orig/path2",
+        ),
+        (
+            "/pyenv/bin:/pyenv/libexec:/orig/path1:/orig/path2",
+            "/orig/path1:/orig/path2",
+        ),
+        (
+            "/pyenv/bin:/some/libexec:/pyenv/hook:/orig/path1:/orig/path2",
+            "/orig/path1:/orig/path2",
+        ),
+        ("/orig/path1:/orig/path2", "/orig/path1:/orig/path2"),
+        (
+            "/orig/path1:/orig/path2:/pyenv/bin:/pyenv/libexec",
+            "/orig/path1:/orig/path2:/pyenv/bin:/pyenv/libexec",
+        ),
+    ],
+)
+def test_fix_env_pyenv(path, orig):
+    env = {"PATH": path, "PYENV_ROOT": "/pyenv"}
+    assert fix_env(env)["PATH"] == orig