Refresh as soon as start

src/client/pythonEnvironments/base/locators/common/nativePythonFinder.ts

@@ -1,19 +1,24 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { Disposable, EventEmitter, Event, Uri, workspace } from 'vscode';
+import { Disposable, EventEmitter, Event, workspace, Uri } from 'vscode';
 import * as ch from 'child_process';
 import * as path from 'path';
 import * as rpc from 'vscode-jsonrpc/node';
 import { PassThrough } from 'stream';
 import { isWindows } from '../../../../common/platform/platformService';
 import { EXTENSION_ROOT_DIR } from '../../../../constants';
 import { traceError, traceInfo, traceLog, traceVerbose, traceWarn } from '../../../../logging';
-import { createDeferred } from '../../../../common/utils/async';
+import { createDeferred, createDeferredFrom } from '../../../../common/utils/async';
 import { DisposableBase, DisposableStore } from '../../../../common/utils/resourceLifecycle';
-import { getPythonSetting } from '../../../common/externalDependencies';
 import { DEFAULT_INTERPRETER_PATH_SETTING_KEY } from '../lowLevel/customWorkspaceLocator';
 import { noop } from '../../../../common/utils/misc';
+import { getConfiguration } from '../../../../common/vscodeApis/workspaceApis';
+import { CONDAPATH_SETTING_KEY } from '../../../common/environmentManagers/conda';
+import { VENVFOLDERS_SETTING_KEY, VENVPATH_SETTING_KEY } from '../lowLevel/customVirtualEnvLocator';
+import { getUserHomeDir } from '../../../../common/utils/platform';
+
+const untildify = require('untildify');
 
 const NATIVE_LOCATOR = isWindows()
     ? path.join(EXTENSION_ROOT_DIR, 'native_locator', 'bin', 'pet.exe')
@@ -43,7 +48,7 @@ export interface NativeEnvManagerInfo {
 
 export interface NativeGlobalPythonFinder extends Disposable {
     resolve(executable: string): Promise<NativeEnvInfo>;
-    refresh(paths: Uri[]): AsyncIterable<NativeEnvInfo>;
+    refresh(): AsyncIterable<NativeEnvInfo>;
 }
 
 interface NativeLog {
@@ -54,9 +59,12 @@ interface NativeLog {
 class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGlobalPythonFinder {
     private readonly connection: rpc.MessageConnection;
 
+    private firstRefreshResults: undefined | (() => AsyncGenerator<NativeEnvInfo, void, unknown>);
+
     constructor() {
         super();
         this.connection = this.start();
+        this.firstRefreshResults = this.refreshFirstTime();
     }
 
     public async resolve(executable: string): Promise<NativeEnvInfo> {
@@ -71,41 +79,82 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
         return environment;
     }
 
-    async *refresh(_paths: Uri[]): AsyncIterable<NativeEnvInfo> {
+    async *refresh(): AsyncIterable<NativeEnvInfo> {
+        if (this.firstRefreshResults) {
+            // If this is the first time we are refreshing,
+            // Then get the results from the first refresh.
+            // Those would have started earlier and cached in memory.
+            const results = this.firstRefreshResults();
+            this.firstRefreshResults = undefined;
+            yield* results;
+        } else {
+            const result = this.doRefresh();
+            let completed = false;
+            void result.completed.finally(() => {
+                completed = true;
+            });
+            const envs: NativeEnvInfo[] = [];
+            let discovered = createDeferred();
+            const disposable = result.discovered((data) => {
+                envs.push(data);
+                discovered.resolve();
+            });
+            do {
+                if (!envs.length) {
+                    await Promise.race([result.completed, discovered.promise]);
+                }
+                if (envs.length) {
+                    const dataToSend = [...envs];
+                    envs.length = 0;
+                    for (const data of dataToSend) {
+                        yield data;
+                    }
+                }
+                if (!completed) {
+                    discovered = createDeferred();
+                }
+            } while (!completed);
+            disposable.dispose();
+        }
+    }
+
+    refreshFirstTime() {
         const result = this.doRefresh();
-        let completed = false;
-        void result.completed.finally(() => {
-            completed = true;
-        });
+        const completed = createDeferredFrom(result.completed);
         const envs: NativeEnvInfo[] = [];
         let discovered = createDeferred();
         const disposable = result.discovered((data) => {
             envs.push(data);
             discovered.resolve();
         });
-        do {
-            if (!envs.length) {
-                await Promise.race([result.completed, discovered.promise]);
-            }
-            if (envs.length) {
-                const dataToSend = [...envs];
-                envs.length = 0;
-                for (const data of dataToSend) {
-                    yield data;
+
+        const iterable = async function* () {
+            do {
+                if (!envs.length) {
+                    await Promise.race([completed.promise, discovered.promise]);
+                }
+                if (envs.length) {
+                    const dataToSend = [...envs];
+                    envs.length = 0;
+                    for (const data of dataToSend) {
+                        yield data;
+                    }
                 }
-            }
-            if (!completed) {
-                discovered = createDeferred();
-            }
-        } while (!completed);
-        disposable.dispose();
+                if (!completed.completed) {
+                    discovered = createDeferred();
+                }
+            } while (!completed.completed);
+            disposable.dispose();
+        };
+
+        return iterable.bind(this);
     }
 
     // eslint-disable-next-line class-methods-use-this
     private start(): rpc.MessageConnection {
         const proc = ch.spawn(NATIVE_LOCATOR, ['server'], { env: process.env });
         const disposables: Disposable[] = [];
-        // jsonrpc package cannot handle messages coming through too quicly.
+        // jsonrpc package cannot handle messages coming through too quickly.
         // Lets handle the messages and close the stream only when
         // we have got the exit event.
         const readable = new PassThrough();
@@ -213,40 +262,86 @@ class NativeGlobalPythonFinderImpl extends DisposableBase implements NativeGloba
                             traceInfo(`Resolved Python Environment ${environment.executable} in ${duration}ms`);
                             discovered.fire(environment);
                         })
-                        .catch((ex) => traceError(`Error in Resolving Python Environment ${data}`, ex));
+                        .catch((ex) => traceError(`Error in Resolving Python Environment ${JSON.stringify(data)}`, ex));
                     trackPromiseAndNotifyOnCompletion(promise);
                 } else {
                     discovered.fire(data);
                 }
             }),
         );
 
-        const pythonPathSettings = (workspace.workspaceFolders || []).map((w) =>
-            getPythonSetting<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY, w.uri.fsPath),
-        );
-        pythonPathSettings.push(getPythonSetting<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY));
-        const pythonSettings = Array.from(new Set(pythonPathSettings.filter((item) => !!item)).values()).map((p) =>
-            // We only want the parent directories.
-            path.dirname(p!),
-        );
         trackPromiseAndNotifyOnCompletion(
-            this.connection
-                .sendRequest<{ duration: number }>('refresh', {
-                    // Send configuration information to the Python finder.
-                    search_paths: (workspace.workspaceFolders || []).map((w) => w.uri.fsPath),
-                    // Also send the python paths that are configured in the settings.
-                    python_path_settings: pythonSettings,
-                    conda_executable: undefined,
-                })
+            this.sendRefreshRequest()
                 .then(({ duration }) => traceInfo(`Native Python Finder completed in ${duration}ms`))
                 .catch((ex) => traceError('Error in Native Python Finder', ex)),
         );
+
         completed.promise.finally(() => disposable.dispose());
         return {
             completed: completed.promise,
             discovered: discovered.event,
         };
     }
+
+    private sendRefreshRequest() {
+        const pythonPathSettings = (workspace.workspaceFolders || []).map((w) =>
+            getPythonSettingAndUntildify<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY, w.uri),
+        );
+        pythonPathSettings.push(getPythonSettingAndUntildify<string>(DEFAULT_INTERPRETER_PATH_SETTING_KEY));
+        // We can have multiple workspaces, each with its own setting.
+        const pythonSettings = Array.from(
+            new Set(
+                pythonPathSettings
+                    .filter((item) => !!item)
+                    // We only want the parent directories.
+                    .map((p) => path.dirname(p!))
+                    /// If setting value is 'python', then `path.dirname('python')` will yield `.`
+                    .filter((item) => item !== '.'),
+            ),
+        );
+
+        return this.connection.sendRequest<{ duration: number }>(
+            'refresh',
+            // Send configuration information to the Python finder.
+            {
+                // This has a special meaning in locator, its lot a low priority
+                // as we treat this as workspace folders that can contain a large number of files.
+                search_paths: (workspace.workspaceFolders || []).map((w) => w.uri.fsPath),
+                // Also send the python paths that are configured in the settings.
+                python_interpreter_paths: pythonSettings,
+                // We do not want to mix this with `search_paths`
+                virtual_env_paths: getCustomVirtualEnvDirs(),
+                conda_executable: getPythonSettingAndUntildify<string>(CONDAPATH_SETTING_KEY),
+                poetry_executable: getPythonSettingAndUntildify<string>('poetryPath'),
+                pipenv_executable: getPythonSettingAndUntildify<string>('pipenvPath'),
+            },
+        );
+    }
+}
+
+/**
+ * Gets all custom virtual environment locations to look for environments.
+ */
+async function getCustomVirtualEnvDirs(): Promise<string[]> {
+    const venvDirs: string[] = [];
+    const venvPath = getPythonSettingAndUntildify<string>(VENVPATH_SETTING_KEY);
+    if (venvPath) {
+        venvDirs.push(untildify(venvPath));
+    }
+    const venvFolders = getPythonSettingAndUntildify<string[]>(VENVFOLDERS_SETTING_KEY) ?? [];
+    const homeDir = getUserHomeDir();
+    if (homeDir) {
+        venvFolders.map((item) => path.join(homeDir, item)).forEach((d) => venvDirs.push(d));
+    }
+    return Array.from(new Set(venvDirs));
+}
+
+function getPythonSettingAndUntildify<T>(name: string, scope?: Uri): T | undefined {
+    const value = getConfiguration('python', scope).get<T>(name);
+    if (typeof value === 'string') {
+        return value ? ((untildify(value as string) as unknown) as T) : undefined;
+    }
+    return value;
 }
 
 export function createNativeGlobalPythonFinder(): NativeGlobalPythonFinder {

src/client/pythonEnvironments/base/locators/lowLevel/nativeLocator.ts

@@ -105,7 +105,7 @@ export class NativeLocator implements ILocator<BasicEnvInfo>, IDisposable {
         const disposables: IDisposable[] = [];
         const disposable = new Disposable(() => disposeAll(disposables));
         this.disposables.push(disposable);
-        for await (const data of this.finder.refresh([])) {
+        for await (const data of this.finder.refresh()) {
             if (data.manager) {
                 switch (toolToKnownEnvironmentTool(data.manager.tool)) {
                     case 'Conda': {