From d65948ca4167434b7e41e9d76687ca008e220af7 Mon Sep 17 00:00:00 2001
From: Alexander Rosenberg <zanderpkg@pm.me>
Date: Sun, 9 Feb 2025 14:05:12 -0800
Subject: [PATCH] Add pyenv support

---
 elisp/eshell-starship.el | 36 +++++++++++++++++++++++++++---------
 init.el                  |  1 +
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/elisp/eshell-starship.el b/elisp/eshell-starship.el
index fa7b442..67877a9 100644
--- a/elisp/eshell-starship.el
+++ b/elisp/eshell-starship.el
@@ -633,9 +633,13 @@ For example, a revert.  If there is no current operation, return nil."
 
 (defun eshell-starship--current-venv ()
   "Return the name of the prompt string for the current venv.
-This requires pyvenv.el to work.  Without it, return nil."
-  (and (bound-and-true-p pyvenv-virtual-env-name)
-       (format " (%s)" pyvenv-virtual-env-name)))
+This requires pyvenv.el or pyenv-mode to work."
+  (concat
+   (and (bound-and-true-p pyvenv-virtual-env-name)
+        (format " (%s)" pyvenv-virtual-env-name))
+   (and (fboundp 'pyenv-mode-version)
+        (when-let ((ver (pyenv-mode-version)))
+          (format " (%s)" ver)))))
 
 (defun eshell-starship--python-status ()
   "Return the prompt string for the python module."
@@ -649,18 +653,32 @@ This requires pyvenv.el to work.  Without it, return nil."
   "The previous `pyvenv-virtual-env' value.
 This does not mean anything if pyvenv.el is not installed.")
 
+(defvar-local eshell-starship--python-last-pyenv nil
+  "The return value of the last `pyenv-mode-version'.
+This does not mean anything if pyenv-mode is not installed.")
+
 (defun eshell-starship--python-postcmd-action ()
   "The postcmd action for the python module."
-  (when (and (boundp 'pyvenv-virtual-env)
-             (not (equal eshell-starship--python-last-pyvenv
-                         pyvenv-virtual-env)))
-    (setq eshell-starship--python-last-pyvenv pyvenv-virtual-env)
-    (eshell-starship-clear-cache-for 'python)))
+  (let ((need-clear nil))
+    (when (and (boundp 'pyvenv-virtual-env)
+               (not (equal eshell-starship--python-last-pyvenv
+                           pyvenv-virtual-env)))
+      (setq eshell-starship--python-last-pyvenv pyvenv-virtual-env
+            need-clear t))
+    (when (fboundp 'pyenv-mode-version)
+      (let ((cur-ver (pyenv-mode-version)))
+        (when (not (equal eshell-starship--python-last-pyenv cur-ver))
+          (setq eshell-starship--python-last-pyenv cur-ver
+                need-clear t))))
+    (when need-clear
+      (eshell-starship-clear-cache-for 'python))))
 
 (eshell-starship-defmodule python
   :extensions '("py" "ipynb")
   :predicate (lambda ()
-               (bound-and-true-p pyvenv-virtual-env))
+               (or (bound-and-true-p pyvenv-virtual-env)
+                   (and (fboundp 'pyenv-mode-version)
+                        (pyenv-mode-version))))
   :files '(".python-version" "Pipfile" "__init__.py" "pyproject.toml"
            "requirements.txt" "setup.py" "tox.ini" "pixi.toml")
   :prefix "via "
diff --git a/init.el b/init.el
index 10f3a20..faab897 100644
--- a/init.el
+++ b/init.el
@@ -1755,6 +1755,7 @@ otherwise, call `bibtex-find-text'."
   :hook (python-ts-mode . trusted-files-eglot-ensure-if-safe))
 ;; python virtual environments
 (use-package pyvenv)
+(use-package pyenv-mode)
 
 ;; java-ts-mode
 (use-package java-ts-mode