Environment info cache class

src/client/pythonEnvironments/base/envsCache.ts

@@ -0,0 +1,91 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { cloneDeep } from 'lodash';
+import { IPersistentState } from '../../common/types';
+import { createGlobalPersistentStore } from '../common/externalDependencies';
+import { PythonEnvInfo } from './info';
+
+/**
+ * Represents the environment info cache to be used by the cache locator.
+ */
+export interface IEnvsCache {
+    /**
+     * Initialization logic to be done outside of the constructor, for example reading from persistent storage.
+     */
+    initialize(): void;
+
+    /**
+     * Return all environment info currently in memory for this session.
+     *
+     * @return An array of cached environment info, or `undefined` if there are none.
+     */
+    getAllEnvs(): PythonEnvInfo[] | undefined;
+
+    /**
+     * Replace all environment info currently in memory for this session.
+     *
+     * @param envs The array of environment info to store in the in-memory cache.
+     */
+    setAllEnvs(envs: PythonEnvInfo[]): void;
+
+    /**
+     * Return a specific environmnent info object.
+     *
+     * @param env The environment info data that will be used to look for an environment info object in the cache.
+     * This object may contain incomplete environment info.
+     * @return The environment info object that matches all non-undefined keys from the `env` param,
+     *  `undefined` otherwise.
+     */
+    getEnv(env: Partial<PythonEnvInfo>): PythonEnvInfo | undefined;
+
+    /**
+     * Writes the content of the in-memory cache to persistent storage.
+     */
+    flush(): Promise<void>;
+}
+
+export type CompleteEnvInfoFunction = (envInfo: PythonEnvInfo) => boolean;
+
+export class PythonEnvInfoCache implements IEnvsCache {
+    private envsList: PythonEnvInfo[] | undefined;
+
+    private persistentStorage: IPersistentState<PythonEnvInfo[]> | undefined;
+
+    constructor(private readonly isComplete: CompleteEnvInfoFunction) {}
+
+    public initialize(): void {
+        this.persistentStorage = createGlobalPersistentStore<PythonEnvInfo[]>('PYTHON_ENV_INFO_CACHE');
+        this.envsList = this.persistentStorage?.value;
+    }
+
+    public getAllEnvs(): PythonEnvInfo[] | undefined {
+        return this.envsList;
+    }
+
+    public setAllEnvs(envs: PythonEnvInfo[]): void {
+        this.envsList = cloneDeep(envs);
+    }
+
+    public getEnv(env: Partial<PythonEnvInfo>): PythonEnvInfo | undefined {
+        // Retrieve all keys with non-undefined values.
+        type EnvParamKeys = keyof typeof env;
+        const keys = (Object.keys(env) as unknown as EnvParamKeys[]).filter((key) => env[key] !== undefined);
+
+        // Return the first object where the values match env's.
+        return this.envsList?.find((info) => {
+            // Check if there is any mismatch between the values of the in-memory info and env.
+            const mismatch = keys.some((key) => info[key] !== env[key]);
+
+            return !mismatch;
+        });
+    }
+
+    public async flush(): Promise<void> {
+        const completeEnvs = this.envsList?.filter(this.isComplete);
+
+        if (completeEnvs?.length) {
+            await this.persistentStorage?.updateValue(completeEnvs);
+        }
+    }
+}

src/client/pythonEnvironments/common/externalDependencies.ts

@@ -4,6 +4,7 @@
 import * as fsapi from 'fs-extra';
 import * as path from 'path';
 import { ExecutionResult, IProcessServiceFactory } from '../../common/process/types';
+import { IPersistentState, IPersistentStateFactory } from '../../common/types';
 import { getOSType, OSType } from '../../common/utils/platform';
 import { IServiceContainer } from '../../ioc/types';
 
@@ -37,3 +38,13 @@ export function arePathsSame(path1: string, path2: string): boolean {
     }
     return path1 === path2;
 }
+
+function getPersistentStateFactory(): IPersistentStateFactory {
+    return internalServiceContainer.get<IPersistentStateFactory>(IPersistentStateFactory);
+}
+
+export function createGlobalPersistentStore<T>(key: string): IPersistentState<T> {
+    const factory = getPersistentStateFactory();
+
+    return factory.createGlobalPersistentState<T>(key, undefined);
+}

src/client/pythonEnvironments/index.ts

@@ -3,6 +3,7 @@
 
 import * as vscode from 'vscode';
 import { IServiceContainer, IServiceManager } from '../ioc/types';
+import { PythonEnvInfoCache } from './base/envsCache';
 import { PythonEnvInfo } from './base/info';
 import { ILocator, IPythonEnvsIterator, PythonLocatorQuery } from './base/locator';
 import { PythonEnvsChangedEvent } from './base/watcher';
@@ -49,13 +50,16 @@ export class PythonEnvironments implements ILocator {
  */
 export function createAPI(): [PythonEnvironments, () => void] {
     const [locators, activateLocators] = initLocators();
+    // Update this to pass in an actual function that checks for env info completeness.
+    const envsCache = new PythonEnvInfoCache(() => true);
 
     return [
         new PythonEnvironments(locators),
         () => {
             activateLocators();
             // Any other activation needed for the API will go here later.
-        }
+            envsCache.initialize();
+        },
     ];
 }
 

src/test/pythonEnvironments/base/envsCache.unit.test.ts

@@ -0,0 +1,137 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import * as assert from 'assert';
+import * as sinon from 'sinon';
+import { CompleteEnvInfoFunction, PythonEnvInfoCache } from '../../../client/pythonEnvironments/base/envsCache';
+import { PythonEnvInfo, PythonEnvKind } from '../../../client/pythonEnvironments/base/info';
+import * as externalDeps from '../../../client/pythonEnvironments/common/externalDependencies';
+
+suite('Environment Info cache', () => {
+    let createGlobalPersistentStoreStub: sinon.SinonStub;
+    let updatedValues: PythonEnvInfo[] | undefined;
+
+    const allEnvsComplete: CompleteEnvInfoFunction = () => true;
+    const envInfoArray = [
+        {
+            id: 'someid1', kind: PythonEnvKind.Conda, name: 'my-conda-env', defaultDisplayName: 'env-one',
+        },
+        {
+            id: 'someid2', kind: PythonEnvKind.Venv, name: 'my-venv-env', defaultDisplayName: 'env-two',
+        },
+        {
+            id: 'someid3', kind: PythonEnvKind.Pyenv, name: 'my-pyenv-env', defaultDisplayName: 'env-three',
+        },
+    ] as PythonEnvInfo[];
+
+    setup(() => {
+        createGlobalPersistentStoreStub = sinon.stub(externalDeps, 'createGlobalPersistentStore');
+        createGlobalPersistentStoreStub.returns({
+            value: envInfoArray,
+            updateValue: async (envs: PythonEnvInfo[]) => {
+                updatedValues = envs;
+                return Promise.resolve();
+            },
+        });
+    });
+
+    teardown(() => {
+        createGlobalPersistentStoreStub.restore();
+        updatedValues = undefined;
+    });
+
+    test('`initialize` reads from persistent storage', () => {
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        envsCache.initialize();
+
+        assert.ok(createGlobalPersistentStoreStub.calledOnce);
+    });
+
+    test('The in-memory env info array is undefined if there is no value in persistent storage when initializing the cache', () => {
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        createGlobalPersistentStoreStub.returns({ value: undefined });
+        envsCache.initialize();
+        const result = envsCache.getAllEnvs();
+
+        assert.strictEqual(result, undefined);
+    });
+
+    test('`getAllEnvs` should return undefined if nothing has been set', () => {
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        const envs = envsCache.getAllEnvs();
+
+        assert.deepStrictEqual(envs, undefined);
+    });
+
+    test('`setAllEnvs` should clone the environment info array passed as a parameter', () => {
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        envsCache.setAllEnvs(envInfoArray);
+        const envs = envsCache.getAllEnvs();
+
+        assert.deepStrictEqual(envs, envInfoArray);
+        assert.strictEqual(envs === envInfoArray, false);
+    });
+
+    test('`getEnv` should return an environment that matches all non-undefined properties of its argument', () => {
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+        envsCache.initialize();
+
+        const result = envsCache.getEnv({ name: 'my-venv-env' });
+
+        assert.deepStrictEqual(result, {
+            id: 'someid2', kind: PythonEnvKind.Venv, name: 'my-venv-env', defaultDisplayName: 'env-two',
+        });
+    });
+
+    test('`getEnv` should return undefined if no environment matches the properties of its argument', () => {
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+        envsCache.initialize();
+
+        const result = envsCache.getEnv({ name: 'my-nonexistent-env' });
+
+        assert.strictEqual(result, undefined);
+    });
+
+    test('`flush` should write complete environment info objects to persistent storage', async () => {
+        const otherEnv = {
+            id: 'someid5',
+            kind: PythonEnvKind.OtherGlobal,
+            name: 'my-other-env',
+            defaultDisplayName: 'env-five',
+        };
+        const updatedEnvInfoArray = [
+            otherEnv, { id: 'someid4', kind: PythonEnvKind.System, name: 'my-system-env' },
+        ] as PythonEnvInfo[];
+        const expected = [
+            otherEnv,
+        ];
+        const envsCache = new PythonEnvInfoCache((env) => env.defaultDisplayName !== undefined);
+
+        envsCache.initialize();
+        envsCache.setAllEnvs(updatedEnvInfoArray);
+        await envsCache.flush();
+
+        assert.deepStrictEqual(updatedValues, expected);
+    });
+
+    test('`flush` should not write to persistent storage if there are no environment info objects in-memory', async () => {
+        const envsCache = new PythonEnvInfoCache((env) => env.kind === PythonEnvKind.MacDefault);
+
+        await envsCache.flush();
+
+        assert.strictEqual(updatedValues, undefined);
+    });
+
+    test('`flush` should not write to persistent storage if there are no complete environment info objects', async () => {
+        const envsCache = new PythonEnvInfoCache((env) => env.kind === PythonEnvKind.MacDefault);
+
+        envsCache.initialize();
+        await envsCache.flush();
+
+        assert.strictEqual(updatedValues, undefined);
+    });
+});