pyenvUtils.ts

import * as fs from 'fs-extra';
import * as path from 'path';
import { Uri } from 'vscode';
import which from 'which';
import {
    EnvironmentManager,
    PythonCommandRunConfiguration,
    PythonEnvironment,
    PythonEnvironmentApi,
    PythonEnvironmentInfo,
} from '../../api';
import { ENVS_EXTENSION_ID } from '../../common/constants';
import { traceError, traceInfo } from '../../common/logging';
import { getWorkspacePersistentState } from '../../common/persistentState';
import { getUserHomeDir, normalizePath, untildify } from '../../common/utils/pathUtils';
import { isWindows } from '../../common/utils/platformUtils';
import {
    isNativeEnvInfo,
    NativeEnvInfo,
    NativeEnvManagerInfo,
    NativePythonEnvironmentKind,
    NativePythonFinder,
} from '../common/nativePythonFinder';
import { shortVersion, sortEnvironments } from '../common/utils';

/**
 * Returns the pyenv root directory from the pyenv executable path.
 * Prefers `PYENV_ROOT` env var when set, otherwise goes up 2 levels from the binary.
 * On POSIX, pyenv binary is at `<root>/bin/pyenv` (e.g. `~/.pyenv/bin/pyenv`).
 * On Windows, pyenv-win binary is at `<root>/bin/pyenv.bat` (e.g. `~/.pyenv/pyenv-win/bin/pyenv.bat`,
 * where `<root>` is `~/.pyenv/pyenv-win`).
 */
export function getPyenvDir(pyenv: string): string {
    const pyenvRoot = process.env.PYENV_ROOT;
    if (pyenvRoot) {
        return pyenvRoot;
    }
    return path.dirname(path.dirname(pyenv));
}

async function findPyenv(): Promise<string | undefined> {
    try {
        return await which('pyenv');
    } catch {
        return undefined;
    }
}

export const PYENV_ENVIRONMENTS = 'Environments';
export const PYENV_VERSIONS = 'Versions';

export const PYENV_PATH_KEY = `${ENVS_EXTENSION_ID}:pyenv:PYENV_PATH`;
export const PYENV_WORKSPACE_KEY = `${ENVS_EXTENSION_ID}:pyenv:WORKSPACE_SELECTED`;
export const PYENV_GLOBAL_KEY = `${ENVS_EXTENSION_ID}:pyenv:GLOBAL_SELECTED`;

let pyenvPath: string | undefined;
export async function clearPyenvCache(): Promise<void> {
    pyenvPath = undefined;
}

async function setPyenv(pyenv: string): Promise<void> {
    pyenvPath = pyenv;
    const state = await getWorkspacePersistentState();
    await state.set(PYENV_PATH_KEY, pyenv);
}

export async function getPyenvForGlobal(): Promise<string | undefined> {
    const state = await getWorkspacePersistentState();
    return await state.get(PYENV_GLOBAL_KEY);
}

export async function setPyenvForGlobal(pyenvPath: string | undefined): Promise<void> {
    const state = await getWorkspacePersistentState();
    await state.set(PYENV_GLOBAL_KEY, pyenvPath);
}

export async function getPyenvForWorkspace(fsPath: string): Promise<string | undefined> {
    const state = await getWorkspacePersistentState();
    const data: { [key: string]: string } | undefined = await state.get(PYENV_WORKSPACE_KEY);
    if (data) {
        try {
            return data[fsPath];
        } catch {
            return undefined;
        }
    }
    return undefined;
}

export async function setPyenvForWorkspace(fsPath: string, envPath: string | undefined): Promise<void> {
    const state = await getWorkspacePersistentState();
    const data: { [key: string]: string } = (await state.get(PYENV_WORKSPACE_KEY)) ?? {};
    if (envPath) {
        data[fsPath] = envPath;
    } else {
        delete data[fsPath];
    }
    await state.set(PYENV_WORKSPACE_KEY, data);
}

export async function setPyenvForWorkspaces(fsPath: string[], envPath: string | undefined): Promise<void> {
    const state = await getWorkspacePersistentState();
    const data: { [key: string]: string } = (await state.get(PYENV_WORKSPACE_KEY)) ?? {};
    fsPath.forEach((s) => {
        if (envPath) {
            data[s] = envPath;
        } else {
            delete data[s];
        }
    });
    await state.set(PYENV_WORKSPACE_KEY, data);
}

export async function getPyenv(): Promise<string | undefined> {
    try {
        if (pyenvPath) {
            if (await fs.exists(untildify(pyenvPath))) {
                return untildify(pyenvPath);
            }
            pyenvPath = undefined;
        }

        const state = await getWorkspacePersistentState();
        const storedPath = await state.get<string>(PYENV_PATH_KEY);
        if (storedPath) {
            if (await fs.exists(untildify(storedPath))) {
                pyenvPath = storedPath;
                traceInfo(`Using pyenv from persistent state: ${pyenvPath}`);
                return untildify(pyenvPath);
            }
            await state.set(PYENV_PATH_KEY, undefined);
        }

        // pyenv-win provides pyenv.bat, not pyenv.exe
        // See: https://github.com/pyenv-win/pyenv-win
        const pyenvBin = isWindows() ? 'pyenv.bat' : 'pyenv';
        const pyenvRoot = process.env.PYENV_ROOT;
        if (pyenvRoot) {
            const pyenvPath = path.join(pyenvRoot, 'bin', pyenvBin);
            if (await fs.exists(pyenvPath)) {
                traceInfo(`Using pyenv from PYENV_ROOT: ${pyenvPath}`);
                return pyenvPath;
            }
        }

        const home = getUserHomeDir();
        if (home) {
            const pyenvPath = path.join(home, '.pyenv', 'bin', pyenvBin);
            if (await fs.exists(pyenvPath)) {
                traceInfo(`Using pyenv from home directory: ${pyenvPath}`);
                return pyenvPath;
            }

            if (isWindows()) {
                const pyenvPathWin = path.join(home, '.pyenv', 'pyenv-win', 'bin', pyenvBin);
                if (await fs.exists(pyenvPathWin)) {
                    traceInfo(`Using pyenv-win from home directory: ${pyenvPathWin}`);
                    return pyenvPathWin;
                }
            }
        }

        pyenvPath = await findPyenv();
        if (pyenvPath) {
            traceInfo(`Using pyenv from PATH: ${pyenvPath}`);
            return pyenvPath;
        }

        traceInfo('pyenv not found in PYENV_ROOT, home directory, or PATH');
        return undefined;
    } catch (ex) {
        const err = ex instanceof Error ? ex : new Error(String(ex));
        (err as Error & { failureStage?: string }).failureStage = `getPyenv`;
        throw err;
    }
}

function nativeToPythonEnv(
    info: NativeEnvInfo,
    api: PythonEnvironmentApi,
    manager: EnvironmentManager,
    pyenv: string,
): PythonEnvironment | undefined {
    if (!(info.prefix && info.executable && info.version)) {
        traceError(`Incomplete pyenv environment info: ${JSON.stringify(info)}`);
        return undefined;
    }

    const pyenvDir = getPyenvDir(pyenv);
    const versionsPath = normalizePath(path.join(pyenvDir, 'versions'));
    const envsPaths = normalizePath(path.join(pyenvDir, 'envs'));
    let group = undefined;
    const normPrefix = normalizePath(info.prefix);
    if (normPrefix.startsWith(versionsPath)) {
        group = PYENV_VERSIONS;
    } else if (normPrefix.startsWith(envsPaths)) {
        group = PYENV_ENVIRONMENTS;
    }

    const sv = shortVersion(info.version);
    const name = info.name || info.displayName || path.basename(info.prefix);
    let displayName = info.displayName || `pyenv (${sv})`;
    if (info.kind === NativePythonEnvironmentKind.pyenvVirtualEnv) {
        displayName = `${name} (${sv})`;
    }

    const shellActivation: Map<string, PythonCommandRunConfiguration[]> = new Map();
    const shellDeactivation: Map<string, PythonCommandRunConfiguration[]> = new Map();

    shellActivation.set('unknown', [{ executable: 'pyenv', args: ['shell', name] }]);
    shellDeactivation.set('unknown', [{ executable: 'pyenv', args: ['shell', '--unset'] }]);

    const environment: PythonEnvironmentInfo = {
        name: name,
        displayName: displayName,
        shortDisplayName: displayName,
        displayPath: info.prefix,
        version: info.version,
        environmentPath: Uri.file(info.prefix),
        description: undefined,
        tooltip: info.prefix,
        execInfo: {
            run: { executable: info.executable },
            shellActivation,
            shellDeactivation,
        },
        sysPrefix: info.prefix,
        group: group,
    };

    return api.createPythonEnvironmentItem(environment, manager);
}

export async function refreshPyenv(
    hardRefresh: boolean,
    nativeFinder: NativePythonFinder,
    api: PythonEnvironmentApi,
    manager: EnvironmentManager,
): Promise<PythonEnvironment[]> {
    traceInfo('Refreshing pyenv environments');
    const data = await nativeFinder.refresh(hardRefresh);

    let pyenv = await getPyenv();

    if (pyenv === undefined) {
        const managers = data
            .filter((e) => !isNativeEnvInfo(e))
            .map((e) => e as NativeEnvManagerInfo)
            .filter((e) => e.tool.toLowerCase() === 'pyenv');

        if (managers.length > 0) {
            pyenv = managers[0].executable;
            await setPyenv(pyenv);
        }
    }

    const envs = data
        .filter((e) => isNativeEnvInfo(e))
        .map((e) => e as NativeEnvInfo)
        .filter(
            (e) =>
                e.kind === NativePythonEnvironmentKind.pyenv || e.kind === NativePythonEnvironmentKind.pyenvVirtualEnv,
        );

    const collection: PythonEnvironment[] = [];

    envs.forEach((e) => {
        if (pyenv) {
            const environment = nativeToPythonEnv(e, api, manager, pyenv);
            if (environment) {
                collection.push(environment);
            }
        }
    });

    traceInfo(`Pyenv discovery complete: ${collection.length} environments found`);
    return sortEnvironments(collection);
}

export async function resolvePyenvPath(
    fsPath: string,
    nativeFinder: NativePythonFinder,
    api: PythonEnvironmentApi,
    manager: EnvironmentManager,
): Promise<PythonEnvironment | undefined> {
    try {
        const e = await nativeFinder.resolve(fsPath);
        if (e.kind !== NativePythonEnvironmentKind.pyenv && e.kind !== NativePythonEnvironmentKind.pyenvVirtualEnv) {
            return undefined;
        }
        const pyenv = await getPyenv();
        if (!pyenv) {
            traceError('Pyenv not found while resolving environment');
            return undefined;
        }

        return nativeToPythonEnv(e, api, manager, pyenv);
    } catch {
        return undefined;
    }
}