Repository: JanDeDobbeleer/oh-my-posh
Pull request #3796 (2023-05-05 15:16:28)

feat(python): add support for pyvenv.cfg prompt
### Prerequisites

- [x] I have read and understood the [contributing guide][CONTRIBUTING.md]
- [x] The commit message follows the [conventional commits][cc] guidelines
- [x] Tests for the changes have been added (for bug fixes/features)
- [ ] Docs have been added/updated (for bug fixes/features)

### Description

Python virtual environments can be created with a user-specified prompt associated to them that overrides the default behaviour of using the environment directory name (see [`venv` module documentation](https://docs.python.org/3/library/venv.html)). This PR implements using that user-specified prompt instead of the environment directory name when provided.

Please advise if the documentation for `.Venv` should be updated (the current description, "the virtual environment name (if present)", seems still appropriate).

### How

If a custom prompt is given on environment creation, it is stored in `pyvenv.cfg`, a file that, according to [PEP 405](https://peps.python.org/pep-0405/) may be found either adjacent to the Python executable or one directory above it (if the executable is a symlink, it is not dereferenced)".

The changes search for a `pyvenv.cfg` file in the relevant locations and for a `prompt` value within that file. If found, the environment directory name resolution is omitted and the retrieved value is used as `.Venv` instead.

[CONTRIBUTING.md]: https://github.com/JanDeDobbeleer/oh-my-posh/blob/main/CONTRIBUTING.md
[cc]: https://www.conventionalcommits.org/en/v1.0.0/#summary



Patch:
diff --git a/src/segments/python.go b/src/segments/python.go
index 6bfb70e1..f1aba4a8 100644
--- a/src/segments/python.go
+++ b/src/segments/python.go
@@ -63,6 +63,10 @@ func (p *Python) loadContext() {
 	if !p.language.props.GetBool(FetchVirtualEnv, true) {
 		return
 	}
+	if prompt := p.pyvenvCfgPrompt(); len(prompt) > 0 {
+		p.Venv = prompt
+		return
+	}
 	venvVars := []string{
 		"VIRTUAL_ENV",
 		"CONDA_ENV_PATH",
@@ -142,3 +146,33 @@ func (p *Python) pyenvVersion() (string, error) {
 	}
 	return parts[0], nil
 }
+
+func (p *Python) pyvenvCfgPrompt() string {
+	path := p.language.env.CommandPath("python")
+	if len(path) == 0 {
+		path = p.language.env.CommandPath("python3")
+	}
+	if len(path) == 0 {
+		return ""
+	}
+	pyvenvDir := filepath.Dir(path)
+	if !p.language.env.HasFilesInDir(pyvenvDir, "pyvenv.cfg") {
+		pyvenvDir = filepath.Dir(pyvenvDir)
+	}
+	if !p.language.env.HasFilesInDir(pyvenvDir, "pyvenv.cfg") {
+		return ""
+	}
+	pyvenvCfg := p.env.FileContent(filepath.Join(pyvenvDir, "pyvenv.cfg"))
+	for _, line := range strings.Split(pyvenvCfg, "\n") {
+		lineSplit := strings.SplitN(line, "=", 2)
+		if len(lineSplit) != 2 {
+			continue
+		}
+		key := strings.TrimSpace(lineSplit[0])
+		if key == "prompt" {
+			value := strings.TrimSpace(lineSplit[1])
+			return value
+		}
+	}
+	return ""
+}


Test patch:
diff --git a/src/segments/python_test.go b/src/segments/python_test.go
index 7be39ba9..7dd72e32 100644
--- a/src/segments/python_test.go
+++ b/src/segments/python_test.go
@@ -2,6 +2,7 @@ package segments
 
 import (
 	"errors"
+	"path/filepath"
 	"testing"
 
 	"github.com/jandedobbeleer/oh-my-posh/src/mock"
@@ -26,6 +27,7 @@ func TestPythonTemplate(t *testing.T) {
 		FetchVersion     bool
 		PythonPath       string
 		ResolveSymlink   ResolveSymlink
+		PyvenvCfg        string
 	}{
 		{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 }}"},
@@ -76,6 +78,15 @@ func TestPythonTemplate(t *testing.T) {
 			Template:       "{{ .Venv }} {{ .Full }}",
 			ResolveSymlink: ResolveSymlink{Path: "/home/user/.pyenv/versions/demo", Err: nil},
 		},
+		{
+			Case:           "pyvenv.cfg prompt",
+			FetchVersion:   true,
+			VirtualEnvName: "VENV",
+			PythonPath:     "/home/user/.pyenv/shims/python",
+			Template:       "{{ if .Venv }}{{ .Venv }} {{ end }}{{ .Major }}.{{ .Minor }}",
+			PyvenvCfg:      "home = /usr/bin/\nprompt = pyvenvCfgPrompt\n",
+			Expected:       "pyvenvCfgPrompt 3.8",
+		},
 	}
 
 	for _, tc := range cases {
@@ -86,6 +97,8 @@ func TestPythonTemplate(t *testing.T) {
 		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("HasFilesInDir", mock2.Anything, "pyvenv.cfg").Return(len(tc.PyvenvCfg) > 0)
+		env.On("FileContent", filepath.Join(filepath.Dir(tc.PythonPath), "pyvenv.cfg")).Return(tc.PyvenvCfg)
 		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)
@@ -122,6 +135,8 @@ func TestPythonPythonInContext(t *testing.T) {
 		env := new(mock.MockedEnvironment)
 		env.On("GOOS").Return("")
 		env.On("PathSeparator").Return("/")
+		env.On("CommandPath", mock2.Anything).Return("")
+		env.On("HasFilesInDir", mock2.Anything, "pyvenv.cfg").Return(false)
 		env.On("Getenv", "VIRTUAL_ENV").Return(tc.VirtualEnvName)
 		env.On("Getenv", "CONDA_ENV_PATH").Return("")
 		env.On("Getenv", "CONDA_DEFAULT_ENV").Return("")