pyenvManager.ts

import * as path from 'path';
import { Disposable, EventEmitter, MarkdownString, ProgressLocation, Uri } from 'vscode';
import {
    DidChangeEnvironmentEventArgs,
    DidChangeEnvironmentsEventArgs,
    EnvironmentChangeKind,
    EnvironmentManager,
    GetEnvironmentScope,
    GetEnvironmentsScope,
    IconPath,
    PythonEnvironment,
    PythonEnvironmentApi,
    PythonProject,
    RefreshEnvironmentsScope,
    ResolveEnvironmentContext,
    SetEnvironmentScope,
} from '../../api';
import { PyenvStrings } from '../../common/localize';
import { traceError, traceInfo } from '../../common/logging';
import { StopWatch } from '../../common/stopWatch';
import { EventNames } from '../../common/telemetry/constants';
import { classifyError } from '../../common/telemetry/errorClassifier';
import { sendTelemetryEvent } from '../../common/telemetry/sender';
import { createDeferred, Deferred } from '../../common/utils/deferred';
import { normalizePath } from '../../common/utils/pathUtils';
import { withProgress } from '../../common/window.apis';
import { PythonProjectManager } from '../../internal.api';
import { getProjectFsPathForScope, tryFastPathGet } from '../common/fastPath';
import { NativePythonFinder } from '../common/nativePythonFinder';
import { getLatest, notifyMissingManagerIfDefault } from '../common/utils';
import {
    clearPyenvCache,
    getPyenv,
    getPyenvForGlobal,
    getPyenvForWorkspace,
    PYENV_VERSIONS,
    refreshPyenv,
    resolvePyenvPath,
    setPyenvForGlobal,
    setPyenvForWorkspace,
    setPyenvForWorkspaces,
} from './pyenvUtils';

export class PyEnvManager implements EnvironmentManager, Disposable {
    private collection: PythonEnvironment[] = [];
    private fsPathToEnv: Map<string, PythonEnvironment> = new Map();
    private globalEnv: PythonEnvironment | undefined;

    private readonly _onDidChangeEnvironment = new EventEmitter<DidChangeEnvironmentEventArgs>();
    public readonly onDidChangeEnvironment = this._onDidChangeEnvironment.event;

    private readonly _onDidChangeEnvironments = new EventEmitter<DidChangeEnvironmentsEventArgs>();
    public readonly onDidChangeEnvironments = this._onDidChangeEnvironments.event;

    constructor(
        private readonly nativeFinder: NativePythonFinder,
        private readonly api: PythonEnvironmentApi,
        private readonly projectManager?: PythonProjectManager,
    ) {
        this.name = 'pyenv';
        this.displayName = 'PyEnv';
        this.preferredPackageManagerId = 'ms-python.python:pip';
        this.tooltip = new MarkdownString(PyenvStrings.pyenvManager, true);
    }

    name: string;
    displayName: string;
    preferredPackageManagerId: string;
    description?: string;
    tooltip: string | MarkdownString;
    iconPath?: IconPath;

    public dispose() {
        this.collection = [];
        this.fsPathToEnv.clear();
    }

    private _initialized: Deferred<void> | undefined;
    async initialize(): Promise<void> {
        if (this._initialized) {
            return this._initialized.promise;
        }
        this._initialized = createDeferred();
        const stopWatch = new StopWatch();
        let result: 'success' | 'tool_not_found' | 'error' = 'success';
        let envCount = 0;
        let toolSource = 'none';
        let errorType: string | undefined;

        try {
            // Check if tool is findable before PET refresh (no settings for pyenv path)
            const preRefreshTool = await getPyenv();
            if (preRefreshTool) {
                toolSource = 'local';
            }

            await withProgress(
                {
                    location: ProgressLocation.Window,
                    title: PyenvStrings.pyenvDiscovering,
                },
                async () => {
                    this.collection = (await refreshPyenv(false, this.nativeFinder, this.api, this)) ?? [];
                    await this.loadEnvMap();

                    this._onDidChangeEnvironments.fire(
                        this.collection.map((e) => ({ environment: e, kind: EnvironmentChangeKind.add })),
                    );
                },
            );

            envCount = this.collection.length;

            // If tool wasn't found via local lookup, check if refresh discovered it via PET
            if (!preRefreshTool) {
                const postRefreshTool = await getPyenv();
                toolSource = postRefreshTool ? 'pet' : 'none';
            }

            if (toolSource === 'none') {
                result = 'tool_not_found';
                if (this.projectManager) {
                    await notifyMissingManagerIfDefault('ms-python.python:pyenv', this.projectManager, this.api);
                }
            }
        } catch (ex) {
            result = 'error';
            errorType = classifyError(ex);
            traceError('Pyenv lazy initialization failed', ex);
        } finally {
            sendTelemetryEvent(EventNames.MANAGER_LAZY_INIT, stopWatch.elapsedTime, {
                managerName: 'pyenv',
                result,
                envCount,
                toolSource,
                errorType,
            });
            this._initialized.resolve();
        }
    }

    async getEnvironments(scope: GetEnvironmentsScope): Promise<PythonEnvironment[]> {
        await this.initialize();

        if (scope === 'all') {
            return Array.from(this.collection);
        }

        if (scope === 'global') {
            return this.collection.filter((env) => env.group === PYENV_VERSIONS);
        }

        if (scope instanceof Uri) {
            const env = this.fromEnvMap(scope);
            if (env) {
                return [env];
            }
        }

        return [];
    }

    async refresh(context: RefreshEnvironmentsScope): Promise<void> {
        if (context === undefined) {
            await withProgress(
                {
                    location: ProgressLocation.Window,
                    title: PyenvStrings.pyenvRefreshing,
                },
                async () => {
                    traceInfo('Refreshing Pyenv Environments');
                    const discard = this.collection.map((c) => c);
                    this.collection = (await refreshPyenv(true, this.nativeFinder, this.api, this)) ?? [];

                    await this.loadEnvMap();

                    const args = [
                        ...discard.map((env) => ({ kind: EnvironmentChangeKind.remove, environment: env })),
                        ...this.collection.map((env) => ({ kind: EnvironmentChangeKind.add, environment: env })),
                    ];

                    this._onDidChangeEnvironments.fire(args);
                },
            );
        }
    }

    async get(scope: GetEnvironmentScope): Promise<PythonEnvironment | undefined> {
        const fastResult = await tryFastPathGet({
            initialized: this._initialized,
            setInitialized: (deferred) => {
                this._initialized = deferred;
            },
            scope,
            label: 'pyenv',
            getProjectFsPath: (s) => getProjectFsPathForScope(this.api, s),
            getPersistedPath: (fsPath) => getPyenvForWorkspace(fsPath),
            resolve: (p) => resolvePyenvPath(p, this.nativeFinder, this.api, this),
            startBackgroundInit: () =>
                withProgress({ location: ProgressLocation.Window, title: PyenvStrings.pyenvDiscovering }, async () => {
                    this.collection = (await refreshPyenv(false, this.nativeFinder, this.api, this)) ?? [];
                    await this.loadEnvMap();
                    this._onDidChangeEnvironments.fire(
                        this.collection.map((e) => ({
                            environment: e,
                            kind: EnvironmentChangeKind.add,
                        })),
                    );
                }),
        });
        if (fastResult) {
            return fastResult.env;
        }

        await this.initialize();
        if (scope instanceof Uri) {
            let env = this.fsPathToEnv.get(normalizePath(scope.fsPath));
            if (env) {
                return env;
            }
            const project = this.api.getPythonProject(scope);
            if (project) {
                env = this.fsPathToEnv.get(normalizePath(project.uri.fsPath));
                if (env) {
                    return env;
                }
            }
        }

        return this.globalEnv;
    }
    async set(scope: SetEnvironmentScope, environment?: PythonEnvironment | undefined): Promise<void> {
        if (scope === undefined) {
            await setPyenvForGlobal(environment?.environmentPath?.fsPath);
        } else if (scope instanceof Uri) {
            const folder = this.api.getPythonProject(scope);
            const fsPath = folder?.uri?.fsPath ?? scope.fsPath;
            if (fsPath) {
                const normalizedFsPath = normalizePath(fsPath);
                if (environment) {
                    this.fsPathToEnv.set(normalizedFsPath, environment);
                } else {
                    this.fsPathToEnv.delete(normalizedFsPath);
                }
                await setPyenvForWorkspace(fsPath, environment?.environmentPath?.fsPath);
            }
        } else if (Array.isArray(scope) && scope.every((u) => u instanceof Uri)) {
            const projects: PythonProject[] = [];
            scope
                .map((s) => this.api.getPythonProject(s))
                .forEach((p) => {
                    if (p) {
                        projects.push(p);
                    }
                });

            const before: Map<string, PythonEnvironment | undefined> = new Map();
            projects.forEach((p) => {
                const normalizedPath = normalizePath(p.uri.fsPath);
                before.set(p.uri.fsPath, this.fsPathToEnv.get(normalizedPath));
                if (environment) {
                    this.fsPathToEnv.set(normalizedPath, environment);
                } else {
                    this.fsPathToEnv.delete(normalizedPath);
                }
            });

            await setPyenvForWorkspaces(
                projects.map((p) => p.uri.fsPath),
                environment?.environmentPath?.fsPath,
            );

            projects.forEach((p) => {
                const b = before.get(p.uri.fsPath);
                if (b?.envId.id !== environment?.envId.id) {
                    this._onDidChangeEnvironment.fire({ uri: p.uri, old: b, new: environment });
                }
            });
        }
    }

    async resolve(context: ResolveEnvironmentContext): Promise<PythonEnvironment | undefined> {
        await this.initialize();

        if (context instanceof Uri) {
            const env = await resolvePyenvPath(context.fsPath, this.nativeFinder, this.api, this);
            if (env) {
                const _collectionEnv = this.findEnvironmentByPath(env.environmentPath.fsPath);
                if (_collectionEnv) {
                    return _collectionEnv;
                }

                this.collection.push(env);
                this._onDidChangeEnvironments.fire([{ kind: EnvironmentChangeKind.add, environment: env }]);

                return env;
            }

            return undefined;
        }
    }

    async clearCache(): Promise<void> {
        await clearPyenvCache();
    }

    private async loadEnvMap() {
        this.globalEnv = undefined;
        this.fsPathToEnv.clear();

        // Try to find a global environment
        const fsPath = await getPyenvForGlobal();

        if (fsPath) {
            this.globalEnv = this.findEnvironmentByPath(fsPath);

            // If the environment is not found, resolve the fsPath. Could be portable conda.
            if (!this.globalEnv) {
                this.globalEnv = await resolvePyenvPath(fsPath, this.nativeFinder, this.api, this);

                // If the environment is resolved, add it to the collection
                if (this.globalEnv) {
                    this.collection.push(this.globalEnv);
                }
            }
        }

        if (!this.globalEnv) {
            this.globalEnv = getLatest(this.collection.filter((e) => e.group === PYENV_VERSIONS));
        }

        // Find any pyenv environments that might be associated with the current projects
        // These are environments whose parent dirs are project dirs.
        const pathSorted = this.collection
            .filter((e) => this.api.getPythonProject(e.environmentPath))
            .sort((a, b) => {
                if (a.environmentPath.fsPath !== b.environmentPath.fsPath) {
                    return a.environmentPath.fsPath.length - b.environmentPath.fsPath.length;
                }
                return a.environmentPath.fsPath.localeCompare(b.environmentPath.fsPath);
            });

        // Try to find workspace environments
        const projects = this.api.getPythonProjects();
        for (const project of projects) {
            const originalPath = project.uri.fsPath;
            const normalizedPath = normalizePath(originalPath);
            const env = await getPyenvForWorkspace(originalPath);

            if (env) {
                const found = this.findEnvironmentByPath(env);

                if (found) {
                    this.fsPathToEnv.set(normalizedPath, found);
                } else {
                    // If not found, resolve the pyenv path. Could be portable pyenv.
                    const resolved = await resolvePyenvPath(env, this.nativeFinder, this.api, this);

                    if (resolved) {
                        // If resolved add it to the collection
                        this.fsPathToEnv.set(normalizedPath, resolved);
                        this.collection.push(resolved);
                    } else {
                        traceError(`Failed to resolve pyenv environment: ${env}`);
                    }
                }
            } else {
                // If there is not an environment already assigned by user to this project
                // then see if there is one in the collection
                if (pathSorted.length === 1) {
                    this.fsPathToEnv.set(normalizedPath, pathSorted[0]);
                } else {
                    // If there is more than one environment then we need to check if the project
                    // is a subfolder of one of the environments
                    const found = pathSorted.find((e) => {
                        const t = this.api.getPythonProject(e.environmentPath)?.uri.fsPath;
                        return t && normalizePath(t) === normalizedPath;
                    });
                    if (found) {
                        this.fsPathToEnv.set(normalizedPath, found);
                    }
                }
            }
        }
    }

    private fromEnvMap(uri: Uri): PythonEnvironment | undefined {
        // Find environment directly using the URI mapping
        const env = this.fsPathToEnv.get(normalizePath(uri.fsPath));
        if (env) {
            return env;
        }

        // Find environment using the Python project for the Uri
        const project = this.api.getPythonProject(uri);
        if (project) {
            return this.fsPathToEnv.get(normalizePath(project.uri.fsPath));
        }

        return undefined;
    }

    private findEnvironmentByPath(fsPath: string): PythonEnvironment | undefined {
        const normalized = normalizePath(fsPath);
        return this.collection.find((e) => {
            const n = normalizePath(e.environmentPath.fsPath);
            return (
                n === normalized ||
                normalizePath(path.dirname(e.environmentPath.fsPath)) === normalized ||
                normalizePath(path.dirname(path.dirname(e.environmentPath.fsPath))) === normalized
            );
        });
    }
}