diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py index 59e9437076..82b2b8306b 100644 --- a/dvc/utils/__init__.py +++ b/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 diff --git a/tests/unit/utils/test_utils.py b/tests/unit/utils/test_utils.py index 3f27a9f061..01ccd6df51 100644 --- a/tests/unit/utils/test_utils.py +++ b/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