micromamba.el/micromamba.el

;;; micromamba.el --- A library for working with micromamba environments -*- lexical-binding: t -*-

;; Copyright (C) 2023 Korytov Pavel

;; Author: Korytov Pavel <thexcloud@gmail.com>
;; Maintainer: Korytov Pavel <thexcloud@gmail.com>
;; Version: 0.1.0
;; Package-Requires: ((emacs "27.1") (pythonic "0.1.0"))
;; Homepage: https://github.com/SqrtMinusOne/micromamba.el
;; Published-At: 2023-06-20

;; This file is NOT part of GNU Emacs.

;; This program is free software: you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.

;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.

;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <https://www.gnu.org/licenses/>.

;;; Commentary:

;; mamba is a reimplementation of the conda package manager in C++.
;; mamba is notably much faster and essentially compatible with conda,
;; so it also works with conda.el.  micromamba, however, implements
;; only a subset of mamba commands, and as such requires a separate
;; integration.
;;
;; The package has two entrypoints:
;; - `micromamba-activate' - activate the environment
;; - `micromamba-deactivate' - deactivate the environment
;;
;; Also see the README at
;; <https://github.com/SqrtMinusOne/micromamba.el> for more
;; information.

;;; Code:
(require 'json)
(require 'pythonic)
(require 'seq)

(defgroup micromamba nil
  "Micromamba (environment manager) integration for Emacs."
  :group 'python)

(defcustom micromamba-executable (executable-find "micromamba")
  "Path to micromamba executable."
  :type 'string
  :group 'micromamba)

(defcustom micromamba-fallback-environment nil
  "An environment that micromamba.el activates by default."
  :type 'string
  :group 'micromamba)

(defcustom micromamba-preactivate-hook nil
  "Hook run before a micromamba environment is activated."
  :type 'hook
  :group 'micromamba)

(defcustom micromamba-postactivate-hook nil
  "Hook run after a micromamba environment is activated."
  :type 'hook
  :group 'micromamba)

(defcustom micromamba-predeactivate-hook nil
  "Hook run before a micromamba environment is deactivated."
  :type 'hook
  :group 'micromamba)

(defcustom micromamba-postdeactivate-hook nil
  "Hook run after a micromamba environment is deactivated."
  :type 'hook
  :group 'micromamba)

(defvar micromamba-env-current-prefix nil
  "Current activated micromamba environment.")

(defvar micromamba--yml-cache nil
  "Stores contents of yaml files.")

(defvar eshell-path-env)

(defun micromamba--call-json (&rest args)
  "Call micromamba and parse the return value as JSON.

Pass ARGS as arguments to the program."
  (unless micromamba-executable
    (user-error "Micromamba-executable is not set!"))
  (with-temp-buffer
    (apply #'call-process micromamba-executable nil t nil args)
    (goto-char (point-min))
    (json-read)))

(defun micromamba-envs ()
  "Get micromamba environments.

Returns an alist with environments, where the key is the name and the
value is the prefix.  Duplicate names are replaced with prefixes."
  (let* ((envs
          (append (alist-get 'envs (micromamba--call-json "env" "list" "--json")) nil))
         (base-env (alist-get 'base\ environment
                              (micromamba--call-json "info" "--json")))
         ;; I have some environments from conda with the same name :-(
         (dupe-name-map (make-hash-table :test 'equal))
         (env-names
          (cl-loop for env in envs
                   for env-name = (if (eq base-env env-name) "base"
                                    (file-name-nondirectory env))
                   do (if (gethash env-name dupe-name-map)
                          (puthash env-name 'dupe dupe-name-map)
                        (puthash env-name t dupe-name-map))
                   collect env-name)))
    (cl-loop for env in envs
             for env-name in env-names
             for name = (if (eq (gethash env-name dupe-name-map) 'dupe)
                            env env-name)
             collect (cons name env))))

(defun micromamba--parse-script-buffer ()
  "Parse bash script buffer generated by micromamba.

E.g. micromamba shell -s bash activate <prefix>.

Returns an alist with the following keys:
- path
- vars-unset
- vars-export
- scripts."
  (let (path vars-unset vars-export scripts)
    (while (not (eobp))
      (cond
       ((looking-at (rx bol "export PATH='"))
        (setq path
              (split-string (buffer-substring-no-properties
                             (+ 13 (point)) (1- (point-at-eol)))
                            ":")))
       ((looking-at (rx bol "unset"))
        (push
         (buffer-substring-no-properties (+ 6 (point)) (point-at-eol))
         vars-unset) )
       ((looking-at (rx bol "export"))
        (save-excursion
          (let ((var-point (+ 7 (point))))
            (re-search-forward (rx "="))
            (let ((var-name (buffer-substring-no-properties var-point (1- (point))))
                  (var-contents (buffer-substring-no-properties
                                 (point) (point-at-eol))))
              (when (string-match-p (rx bos "'" (* nonl) "'" eos) var-contents)
                (setq var-contents
                      (substring var-contents 1 (1- (length var-contents)))))
              (push (cons var-name var-contents) vars-export)))))
       ((looking-at (rx bol ". \""))
        (push (buffer-substring-no-properties (+ 4 (point))
                                              (1- (point-at-eol)))
              scripts)))
      (forward-line))
    `((path . ,path)
      (vars-unset . ,vars-unset)
      (vars-export . ,vars-export)
      (scripts . ,scripts))))

(defun micromamba--find-env-yml (dir) ;; adapted from conda.el
  "Find an environment.yml or .yaml in DIR or its parent directories."
  (let ((containing-path (locate-dominating-file dir
                          (lambda (parent)
                            (directory-files parent nil "environment.[yml|yaml]")))))
    (when containing-path
      (let ((yml-candidate
             (concat (file-name-as-directory containing-path) "environment.yml"))
            (yaml-candidate
             (concat (file-name-as-directory containing-path) "environment.yaml")))
        (or (when (file-readable-p yml-candidate) yml-candidate)
            (when (file-readable-p yaml-candidate) yaml-candidate))))))

(defun micromamba--read-file-into-string (filename)
  "Read the contents of FILENAME into a string."
  (with-temp-buffer
    (let ((coding-system-for-read 'utf-8))
      (insert-file-contents filename)
      (buffer-string))))

(defun micromamba--get-name-from-env-yml-contents (env-yml-contents)
  "Pull the `name` property out of a stringified YAML file"
  (save-match-data
    (when (string-match "name:[ ]*\\([A-z0-9-_.]+\\)[ ]*$" env-yml-contents)
      (match-string 1 env-yml-contents))))

(defun micromamba--get-name-from-env-yml (filename) ;; adapted from conda.el
  "Pull the `name` property out of the YAML file at FILENAME."
  (when filename
    (let ((filename (file-truename filename)))
      (or (unless (file-has-changed-p filename))
          (cdr (assoc filename micromamba--yml-cache))
          (setf (alist-get filename micromamba--yml-cache)
                (micromamba--get-name-from-env-yml-contents
                 (micromamba--read-file-into-string filename)))))))

(defun micromamba--infer-env-from-buffer () ;; adapted from conda.el
  "Search up the project tree for an `environment.yml` defining a conda env.

Return `micromamba-fallback-environment' if not found."
  (let* ((filename (buffer-file-name))
         (working-dir (if filename
                          (file-name-directory filename)
                        default-directory)))
    (when working-dir
      (or
       (micromamba--get-name-from-env-yml (micromamba--find-env-yml working-dir))
       micromamba-fallback-environment))))

(defun micromamba--get-activation-parameters (prefix)
  "Get activation parameters for the environment PREFIX.

The parameters value is an alist as defined by
`micromamba--parse-script-buffer'."
  (with-temp-buffer
    (call-process micromamba-executable nil t nil
                  "shell" "activate" prefix "-s" "bash")
    (goto-char (point-min))
    (micromamba--parse-script-buffer)))

(defun micromamba--get-deactivation-parameters ()
  "Get deactivation parameters for the current evironment.

The parameters value is an alist as defined by
`micromamba--parse-script-buffer'."
  (with-temp-buffer
    (call-process micromamba-executable nil t nil
                  "shell" "deactivate" "-s" "bash")
    (goto-char (point-min))
    (micromamba--parse-script-buffer)))

(defun micromamba--apply-env (parameters)
  "Apply PARAMETERS to the current environment.

The parameters value is an alist as defined by
`micromamba--parse-script-buffer'."
  (unless (alist-get 'path parameters)
    (user-error "Something went wrong.  Cannot get PATH"))
  (setq exec-path (alist-get 'path parameters))
  (setenv "PATH" (string-join (alist-get 'path parameters) ":"))
  (dolist (var-name (alist-get 'vars-unset parameters))
    (setenv var-name nil))
  (dolist (var (alist-get 'vars-export parameters))
    (setenv (car var) (cdr var)))
  (setq eshell-path-env (getenv "PATH")))

(defun micromamba--env-name-to-prefix (name)
  "Get the prefix of an environment with NAME."
  (let ((envs (micromamba-envs)))
    (alist-get name envs nil nil #'equal)))

;; "public" functions
;;;###autoload
(defun micromamba-activate (prefix)
  "Switch to environment with PREFIX (path).  Prompt if called interactively.

If some environments have duplicate names, these names are replaced by
full paths."
  (interactive
   (list (let ((envs (micromamba-envs)))
           (alist-get
            (completing-read "Choose a micromamba environment: " envs
                             nil t)
            envs nil nil #'equal))))
  ;; To allow calling the function with env name as well
  (unless (string-match-p (rx bos "/") prefix)
    (setq prefix (micromamba--env-name-to-prefix prefix))
    (unless prefix
      (user-error "Environment %s not found" prefix)))
  (micromamba-deactivate)
  (setq micromamba-env-current-prefix prefix)
  (run-hooks 'micromamba-preactivate-hook)
  ;; conda.el was doing that, so why not?
  (pythonic-activate prefix)
  (setq python-shell-virtualenv-root prefix)
  (micromamba--apply-env
   (micromamba--get-activation-parameters prefix))
  (run-hooks 'micromamba-postactivate-hook)
  (message "Switched to micromamba environment: %s" prefix))

;;;###autoload
(defun micromamba-deactivate ()
  "Deactivate the current environment."
  (interactive)
  (when (bound-and-true-p micromamba-env-current-prefix)
    (run-hooks 'micromamba-predeactivate-hook)
    (setq python-shell-virtualenv-root nil)
    (micromamba--apply-env
     (micromamba--get-deactivation-parameters))
    (setq micromamba-env-current-prefix nil)
    (run-hooks 'micromamba-postdeactivate-hook)))

;;;###autoload
(defun micromamba-env-activate-for-buffer ()
  "Activate the conda environment implied by the current buffer.

This can be set by a buffer-local or project-local variable (e.g. a
`.dir-locals.el` that defines `conda-project-env-path`), or inferred from an
`environment.yml` or similar at the project level."
  (interactive)
  (let ((inferred-env (micromamba--infer-env-from-buffer)))
    (when (and inferred-env
               (not (string= micromamba-env-current-prefix
                             (micromamba--env-name-to-prefix inferred-env))))
      (micromamba-activate inferred-env))))

(defun micromamba--switch-buffer-auto-activate (&rest args)
  "Add Conda environment activation if a buffer has a file, handling ARGS."
  (let ((filename (buffer-file-name)))
    (when filename
      (with-demoted-errors "Error: %S"
        (micromamba-env-activate-for-buffer)))))

;;;###autoload
(define-minor-mode micromamba-env-autoactivate-mode
  "Toggle conda-env-autoactivate mode.

This mode automatically tries to activate a conda environment for the current
buffer."
  ;; The initial value.
  :init-value nil
  ;; The indicator for the mode line.
  :lighter nil
  ;; The minor mode bindings.
  :keymap nil
  ;; Kwargs
  :group 'micromamba
  :global t
  ;; Forms
  (if micromamba-env-autoactivate-mode ;; already on, now switching off
    (add-hook 'window-selection-change-functions
                 #'micromamba--switch-buffer-auto-activate)
    (remove-hook 'window-selection-change-functions #'micromamba--switch-buffer-auto-activate)))

(provide 'micromamba)
;;; micromamba.el ends here