From 132de572af58cae15a742c0d5bfd365e0effb4f1 Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <ruslan@iterative.ai>
Date: Wed, 2 Oct 2019 06:02:40 +0300
Subject: [PATCH] run: improve pyenv PATH modification criteria

https://discordapp.com/channels/485586884165107732/485596304961962003/628766062472724510

Related to #2506
---
 dvc/utils/__init__.py          | 39 +++++++++++++++++++++++++++-------
 tests/unit/utils/test_utils.py | 21 +++++++++++++++---
 2 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py
index 82b2b8306..cc5e3bd78 100644
--- a/dvc/utils/__init__.py
+++ b/dvc/utils/__init__.py
@@ -272,22 +272,45 @@ def fix_env(env=None):
     #
     # where
     #
-    #     PYENV_BIN_PATH - starts with $PYENV_ROOT, see pyenv-exec source code.
+    #     PYENV_BIN_PATH - might not start with $PYENV_ROOT if we are running
+    #         `system` version of the command, see pyenv-exec source code.
     #     bin_path - might not start with $PYENV_ROOT as it runs realpath on
-    #         it, see pyenv source code.
+    #         it, but always has `libexec` part in it, see pyenv source code.
     #     plugin_bin - might contain more than 1 entry, which start with
     #         $PYENV_ROOT, see pyenv source code.
     #
+    # Also, we know that whenever pyenv is running, it exports these env vars:
+    #
+    #     PYENV_DIR
+    #     PYENV_HOOK_PATH
+    #     PYENV_VERSION
+    #     PYENV_ROOT
+    #
     # 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")
+    path = env.get("PATH", "")
+    parts = path.split(":")
+    bin_path = parts[1] if len(parts) > 2 else ""
+    pyenv_dir = env.get("PYENV_DIR")
+    pyenv_hook_path = env.get("PYENV_HOOK_PATH")
+    pyenv_version = env.get("PYENV_VERSION")
     pyenv_root = env.get("PYENV_ROOT")
-    if path and pyenv_root and path.startswith(pyenv_root):
+
+    env_matches = all([pyenv_dir, pyenv_hook_path, pyenv_version, pyenv_root])
+
+    bin_path_matches = os.path.basename(bin_path) == "libexec"
+
+    # NOTE: we don't support pyenv-win
+    if os.name != "nt" and env_matches and bin_path_matches:
         # 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]
+        parts = parts[2:]
+
+        if parts:
+            # removing plugin_bin from the left
+            plugin_bin = os.path.join(pyenv_root, "plugins")
+            while parts[0].startswith(plugin_bin):
+                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 01ccd6df5..5678994d4 100644
--- a/tests/unit/utils/test_utils.py
+++ b/tests/unit/utils/test_utils.py
@@ -1,3 +1,5 @@
+import os
+
 import pytest
 
 from dvc.utils import to_chunks, fix_env
@@ -30,11 +32,15 @@ def test_to_chunks_num_chunks(num_chunks, expected_chunks):
     assert result == expected_chunks
 
 
+@pytest.mark.skipif(os.name == "nt", reason="pyenv-win is not supported")
 @pytest.mark.parametrize(
     "path, orig",
     [
         (
-            "/pyenv/bin:/pyenv/libexec:/pyenv/hook:/orig/path1:/orig/path2",
+            (
+                "/pyenv/bin:/pyenv/libexec:/pyenv/plugins/plugin:"
+                "/orig/path1:/orig/path2"
+            ),
             "/orig/path1:/orig/path2",
         ),
         (
@@ -42,7 +48,10 @@ def test_to_chunks_num_chunks(num_chunks, expected_chunks):
             "/orig/path1:/orig/path2",
         ),
         (
-            "/pyenv/bin:/some/libexec:/pyenv/hook:/orig/path1:/orig/path2",
+            (
+                "/pyenv/bin:/some/libexec:/pyenv/plugins/plugin:"
+                "/orig/path1:/orig/path2"
+            ),
             "/orig/path1:/orig/path2",
         ),
         ("/orig/path1:/orig/path2", "/orig/path1:/orig/path2"),
@@ -53,5 +62,11 @@ def test_to_chunks_num_chunks(num_chunks, expected_chunks):
     ],
 )
 def test_fix_env_pyenv(path, orig):
-    env = {"PATH": path, "PYENV_ROOT": "/pyenv"}
+    env = {
+        "PATH": path,
+        "PYENV_ROOT": "/pyenv",
+        "PYENV_VERSION": "3.7.2",
+        "PYENV_DIR": "/some/dir",
+        "PYENV_HOOK_PATH": "/some/hook/path",
+    }
     assert fix_env(env)["PATH"] == orig