Repository: JanDeDobbeleer/oh-my-posh
Pull request #2122 (2022-04-21 16:42:50)

feat(python): real pyenv support
If `python` points at a pyenv shim, use `pyenv version-name` to
determine Python version and Venv information.

### Prerequisites

- [x] I have read and understood the `CONTRIBUTING` guide
- [x] The commit message follows the [conventional commits][cc] guidelines
- [x] Tests for the changes have been added (for bug fixes/features)
- [x] Docs have been added / updated (for bug fixes/features)

### Description


bewing: This approach shaves about 100msec off python prompt generation on my WSL environment (since we never run `python` directly, and `pyenv` is 100% bash)

debug log for the segment:
```
python(true)         -  51 ms -   virt 3.9.9
2022/04/21 11:43:37 debug: Pwd
/home/bewing/python
2022/04/21 11:43:37 Pwd duration: 870ns, args:
2022/04/21 11:43:37 PathSeparator duration: 62ns, args:
2022/04/21 11:43:37 debug: HasFiles
true
2022/04/21 11:43:37 HasFiles duration: 14.724µs, args: *.py
2022/04/21 11:43:37 debug: RunCommand
virt
2022/04/21 11:43:37 RunCommand duration: 35.492697ms, args: pyenv version-name
2022/04/21 11:43:37 TemplateCache duration: 524ns, args:
2022/04/21 11:43:37 TemplateCache duration: 938ns, args:
```
bewing: Need to refactor this code to make it mockable/testable.  Changes to `environment` probably needed
bewing: Think code and tests are good now.  Will look at docs next.  Do we want to feature gate this, or no?
JanDeDobbeleer: @bewing I wouldn't gate it, just ship. If you tested locally, we'll find out just how representative it was 😃
bewing: Bah, I thought I was testing locally, but I wasn't.  Not ready for review yet, but I'm not sure how to mark it as draft again
bewing: Found the error in my test case, fixed it.  Really really think it's ready now.

Patch:
diff --git a/docs/docs/segments/python.md b/docs/docs/segments/python.md
index fe79e9a1..8331af12 100644
--- a/docs/docs/segments/python.md
+++ b/docs/docs/segments/python.md
@@ -7,7 +7,7 @@ sidebar_label: Python
 ## What
 
 Display the currently active python version and virtualenv.
-Supports conda, virtualenv and pyenv.
+Supports conda, virtualenv and pyenv (if python points to pyenv shim).
 
 ## Sample Configuration
 
@@ -26,7 +26,6 @@ Supports conda, virtualenv and pyenv.
 
 - home_enabled: `boolean` - display the segment in the HOME folder or not - defaults to `false`
 - fetch_virtual_env: `boolean` - fetch the name of the virtualenv or not - defaults to `true`
-- use_python_version_file: `boolean` - Use pyenv `.python-version` files to determine virtual env - defaults to `false`
 - display_default: `boolean` - show the name of the virtualenv when it's default (`system`, `base`)
 or not - defaults to `true`
 - fetch_version: `boolean` - fetch the python version - defaults to `true`
diff --git a/src/environment/shell.go b/src/environment/shell.go
index 4a09f080..15b9293f 100644
--- a/src/environment/shell.go
+++ b/src/environment/shell.go
@@ -139,7 +139,9 @@ type Environment interface {
 	HasFolder(folder string) bool
 	HasParentFilePath(path string) (fileInfo *FileInfo, err error)
 	HasFileInParentDirs(pattern string, depth uint) bool
+	ResolveSymlink(path string) (string, error)
 	DirMatchesOneOf(dir string, regexes []string) bool
+	CommandPath(command string) string
 	HasCommand(command string) bool
 	FileContent(file string) string
 	LsDir(path string) []fs.DirEntry
@@ -401,6 +403,10 @@ func (env *ShellEnvironment) HasFolder(folder string) bool {
 	return hasFolder
 }
 
+func (env *ShellEnvironment) ResolveSymlink(path string) (string, error) {
+	return filepath.EvalSymlinks(path)
+}
+
 func (env *ShellEnvironment) FileContent(file string) string {
 	defer env.trace(time.Now(), "FileContent", file)
 	if !filepath.IsAbs(file) {
@@ -500,22 +506,29 @@ func (env *ShellEnvironment) RunShellCommand(shell, command string) string {
 	return ""
 }
 
-func (env *ShellEnvironment) HasCommand(command string) bool {
+func (env *ShellEnvironment) CommandPath(command string) string {
 	defer env.trace(time.Now(), "HasCommand", command)
-	if _, ok := env.cmdCache.get(command); ok {
-		return true
+	if path, ok := env.cmdCache.get(command); ok {
+		return path
 	}
 	path, err := exec.LookPath(command)
 	if err == nil {
 		env.cmdCache.set(command, path)
-		return true
+		return path
 	}
 	path, err = env.LookWinAppPath(command)
 	if err == nil {
 		env.cmdCache.set(command, path)
+		return path
+	}
+	env.log(Error, "CommandPath", err.Error())
+	return ""
+}
+
+func (env *ShellEnvironment) HasCommand(command string) bool {
+	if path := env.CommandPath(command); path != "" {
 		return true
 	}
-	env.log(Error, "HasCommand", err.Error())
 	return false
 }
 
diff --git a/src/mock/environment.go b/src/mock/environment.go
index 589b8a69..2aa15c52 100644
--- a/src/mock/environment.go
+++ b/src/mock/environment.go
@@ -42,6 +42,11 @@ func (env *MockedEnvironment) HasFolder(folder string) bool {
 	return args.Bool(0)
 }
 
+func (env *MockedEnvironment) ResolveSymlink(path string) (string, error) {
+	args := env.Called(path)
+	return args.String(0), args.Error(1)
+}
+
 func (env *MockedEnvironment) FileContent(file string) string {
 	args := env.Called(file)
 	return args.String(0)
@@ -77,6 +82,11 @@ func (env *MockedEnvironment) Platform() string {
 	return args.String(0)
 }
 
+func (env *MockedEnvironment) CommandPath(command string) string {
+	args := env.Called(command)
+	return args.String(0)
+}
+
 func (env *MockedEnvironment) HasCommand(command string) bool {
 	args := env.Called(command)
 	return args.Bool(0)
diff --git a/src/segments/python.go b/src/segments/python.go
index 4fc6ea6e..66c52ece 100644
--- a/src/segments/python.go
+++ b/src/segments/python.go
@@ -3,7 +3,7 @@ package segments
 import (
 	"oh-my-posh/environment"
 	"oh-my-posh/properties"
-	"oh-my-posh/regex"
+	"path/filepath"
 	"strings"
 )
 
@@ -31,6 +31,10 @@ func (p *Python) Init(props properties.Properties, env environment.Environment)
 		loadContext: p.loadContext,
 		inContext:   p.inContext,
 		commands: []*cmd{
+			{
+				getVersion: p.pyenvVersion,
+				regex:      `(?P<version>((?P<major>[0-9]+).(?P<minor>[0-9]+).(?P<patch>[0-9]+)))`,
+			},
 			{
 				executable: "python",
 				args:       []string{"--version"},
@@ -60,7 +64,6 @@ func (p *Python) loadContext() {
 		"VIRTUAL_ENV",
 		"CONDA_ENV_PATH",
 		"CONDA_DEFAULT_ENV",
-		"PYENV_VERSION",
 	}
 	var venv string
 	for _, venvVar := range venvVars {
@@ -71,15 +74,6 @@ func (p *Python) loadContext() {
 			break
 		}
 	}
-	if !p.language.props.GetBool(UsePythonVersionFile, false) {
-		return
-	}
-	if f, err := p.language.env.HasParentFilePath(".python-version"); err == nil {
-		contents := strings.Split(p.language.env.FileContent(f.Path), "\n")
-		if contents[0] != "" && regex.MatchString("[0-9]+.[0-9]+.[0-9]", contents[0]) == false && p.canUseVenvName(contents[0]) {
-			p.Venv = contents[0]
-		}
-	}
 }
 
 func (p *Python) inContext() bool {
@@ -101,3 +95,48 @@ func (p *Python) canUseVenvName(name string) bool {
 	}
 	return true
 }
+
+func (p *Python) pyenvVersion() (string, error) {
+	// Use `pyenv root` instead of $PYENV_ROOT?
+	// Is our Python executable at $PYENV_ROOT/bin/python ?
+	// Should p.env expose command paths?
+	path := p.env.CommandPath("python")
+	if path == "" {
+		path = p.env.CommandPath("python3")
+	}
+	if path == "" {
+		return "", nil
+	}
+	// TODO:  pyenv-win has this at $PYENV_ROOT/pyenv-win/shims
+	if path != filepath.Join(p.env.Getenv("PYENV_ROOT"), "shims", "python") {
+		return "", nil
+	}
+	// pyenv version-name will return current version or virtualenv
+	cmdOutput, err := p.env.RunCommand("pyenv", "version-name")
+	if err != nil {
+		// TODO: Improve reporting
+		return "", nil
+	}
+	versionString := strings.Split(cmdOutput, ":")[0]
+	if versionString == "" {
+		return "", nil
+	}
+
+	// $PYENV_ROOT/versions + versionString (symlinks resolved) == $PYENV_ROOT/versions/(version)[/envs/(virtualenv)]
+	realPath, err := p.env.ResolveSymlink(filepath.Join(p.env.Getenv("PYENV_ROOT"), "versions", versionString))
+	if err != nil {
+		return "", nil
+	}
+	// ../versions/(version)[/envs/(virtualenv)]
+	shortPath, err := filepath.Rel(filepath.Join(p.env.Getenv("PYENV_ROOT"), "versions"), realPath)
+	if err != nil {
+		return "", nil
+	}
+	// Unset whatever loadContext thinks Venv should be
+	p.Venv = ""
+	parts := strings.Split(shortPath, string(filepath.Separator))
+	if len(parts) > 2 && p.canUseVenvName(parts[2]) {
+		p.Venv = parts[2]
+	}
+	return parts[0], nil
+}


Test patch:
diff --git a/src/segments/python_test.go b/src/segments/python_test.go
index 62066c5f..a8ba42f9 100644
--- a/src/segments/python_test.go
+++ b/src/segments/python_test.go
@@ -11,6 +11,10 @@ import (
 )
 
 func TestPythonTemplate(t *testing.T) {
+	type ResolveSymlink struct {
+		Path string
+		Err  error
+	}
 	cases := []struct {
 		Case             string
 		Expected         string
@@ -18,7 +22,8 @@ func TestPythonTemplate(t *testing.T) {
 		Template         string
 		VirtualEnvName   string
 		FetchVersion     bool
-		PyenvLocal       string
+		PythonPath       string
+		ResolveSymlink   ResolveSymlink
 	}{
 		{Case: "No virtual env present", FetchVersion: true, Expected: "3.8.4", Template: "{{ if .Venv }}{{ .Venv }} {{ end }}{{ .Full }}"},
 		{Case: "Virtual env present", FetchVersion: true, Expected: "VENV 3.8.4", VirtualEnvName: "VENV", Template: "{{ if .Venv }}{{ .Venv }} {{ end }}{{ .Full }}"},
@@ -44,39 +49,39 @@ func TestPythonTemplate(t *testing.T) {
 			Template:       "{{ if ne .Venv \"default\" }}{{ .Venv }} {{ end }}{{ .Major }}.{{ .Minor }}",
 		},
 		{
-			Case:         "Pyenv show env",
-			FetchVersion: true,
-			Expected:     "VENV 3.8",
-			PyenvLocal:   "VENV\n",
-			Template:     "{{ if ne .Venv \"default\" }}{{ .Venv }} {{ end }}{{ .Major }}.{{ .Minor }}",
+			Case:           "Pyenv show env",
+			FetchVersion:   true,
+			Expected:       "VENV 3.8",
+			PythonPath:     "/home/user/.pyenv/shims/python",
+			VirtualEnvName: "VENV",
+			Template:       "{{ if ne .Venv \"default\" }}{{ .Venv }} {{ end }}{{ .Major }}.{{ .Minor }}",
+			ResolveSymlink: ResolveSymlink{Path: "/home/user/.pyenv/versions/3.8.8/envs/VENV", Err: nil},
 		},
 		{
-			Case:         "Pyenv skip version",
-			FetchVersion: true,
-			Expected:     "3.8",
-			PyenvLocal:   "3.8.7\n",
-			Template:     "{{ if ne .Venv \"default\" }}{{ .Venv }} {{ end }}{{ .Major }}.{{ .Minor }}",
+			Case:           "Pyenv no venv",
+			FetchVersion:   true,
+			Expected:       "3.8",
+			PythonPath:     "/home/user/.pyenv/shims/python",
+			Template:       "{{ if ne .Venv \"default\" }}{{ .Venv }} {{ end }}{{ .Major }}.{{ .Minor }}",
+			ResolveSymlink: ResolveSymlink{Path: "/home/user.pyenv/versions/3.8.8", Err: nil},
 		},
 	}
 
 	for _, tc := range cases {
 		env := new(mock.MockedEnvironment)
 		env.On("HasCommand", "python").Return(true)
+		env.On("CommandPath", "mock.Anything").Return(tc.PythonPath)
 		env.On("RunCommand", "python", []string{"--version"}).Return("Python 3.8.4", nil)
+		env.On("RunCommand", "pyenv", []string{"version-name"}).Return(tc.VirtualEnvName, nil)
 		env.On("HasFiles", "*.py").Return(true)
-		env.On("HasParentFilePath", ".python-version").Return(&environment.FileInfo{
-			ParentFolder: "/usr/home/project",
-			Path:         "/usr/home/project/.python-version",
-			IsDir:        false,
-		}, nil)
-		env.On("FileContent", "/usr/home/project/.python-version").Return(tc.PyenvLocal)
 		env.On("Getenv", "VIRTUAL_ENV").Return(tc.VirtualEnvName)
 		env.On("Getenv", "CONDA_ENV_PATH").Return(tc.VirtualEnvName)
 		env.On("Getenv", "CONDA_DEFAULT_ENV").Return(tc.VirtualEnvName)
-		env.On("Getenv", "PYENV_VERSION").Return(tc.VirtualEnvName)
+		env.On("Getenv", "PYENV_ROOT").Return("/home/user/.pyenv")
 		env.On("PathSeparator").Return("")
 		env.On("Pwd").Return("/usr/home/project")
 		env.On("Home").Return("/usr/home")
+		env.On("ResolveSymlink", "mock.Anything").Return(tc.ResolveSymlink.Path, tc.ResolveSymlink.Err)
 		props := properties.Map{
 			properties.FetchVersion: tc.FetchVersion,
 			UsePythonVersionFile:    true,