Make PythonLocatorQuery.searchLocations more explicit about the non-r…

…ooted cases.
Eric Snow (ericsnowcurrently) · Sep 29, 2020 · Sep 22, 2020 · Sep 23, 2020 · Sep 23, 2020 · Sep 21, 2020
commit ce53e541fbaf819dcadb90cd99c633273abd55f5
diff --git a/src/client/pythonEnvironments/base/locator.ts b/src/client/pythonEnvironments/base/locator.ts
@@ -82,19 +82,31 @@ export type BasicPythonLocatorQuery = {
     kinds?: PythonEnvKind[];
 };
 
+/**
+ * The portion of a query related to env search locations.
+ */
+export type SearchLocations = {
+    /**
+     * The locations under which to look for environments.
+     */
+    roots: Uri[];
+    /**
+     * If true, also look for environments that do not have a search location.
+     */
+    includeNonRooted?: boolean;
+};
+
 /**
  * The full set of possible info to send to a locator when requesting environments.
  *
  * This is directly correlated with the `PythonEnvsChangedEvent`
  * emitted by watchers.
- *
- * @prop - searchLocations - if provided, results should be limited to
- *         within these locations; `null` means "include envs that have
- *         no search location"; if not provided, the search location
- *         of each enviroment is not considered when filtering
  */
 export type PythonLocatorQuery = BasicPythonLocatorQuery & {
-    searchLocations?: (Uri | null)[] | null;
+    /**
+     * If provided, results should be limited to within these locations.
+     */
+    searchLocations?: SearchLocations;
 };
 
 type QueryForEvent<E> = E extends PythonEnvsChangedEvent ? PythonLocatorQuery : BasicPythonLocatorQuery;

diff --git a/src/client/pythonEnvironments/base/locatorUtils.ts b/src/client/pythonEnvironments/base/locatorUtils.ts
@@ -18,18 +18,16 @@ export function getQueryFilter(query: PythonLocatorQuery): (env: PythonEnvInfo)
     const kinds = (query.kinds !== undefined && query.kinds.length > 0)
         ? query.kinds
         : undefined;
-    let includeGlobal = !query.searchLocations || query.searchLocations.length === 0;
-    let locationFilters: ((u: Uri) => boolean)[] | undefined;
-    if (!includeGlobal) {
-        const candidates = query.searchLocations!.filter((u) => !!u);
-        includeGlobal = candidates.length < query.searchLocations!.length;
-        if (candidates.length > 0) {
-            locationFilters = candidates.map((loc) => getURIFilter(loc!, {
-                checkParent: true,
-                checkExact: true,
-            }));
+    let includeNonRooted = true;
+    if (query.searchLocations !== undefined) {
+        if (query.searchLocations.includeNonRooted !== undefined) {
+            includeNonRooted = query.searchLocations.includeNonRooted;
+        } else {
+            // We default to `false`.
+            includeNonRooted = false;
         }
     }
+    const locationFilters = getSearchLocationFilters(query);
     function checkKind(env: PythonEnvInfo): boolean {
         if (kinds === undefined) {
             return true;
@@ -39,20 +37,15 @@ export function getQueryFilter(query: PythonLocatorQuery): (env: PythonEnvInfo)
     function checkSearchLocation(env: PythonEnvInfo): boolean {
         if (env.searchLocation === undefined) {
             // It is not a "rooted" env.
-            return includeGlobal;
-        }
-
-        // It is a "rooted" env.
-        if (locationFilters === undefined) {
-            if (query.searchLocations === null) {
-                // The caller explicitly refused rooted envs.
-                return false;
+            return includeNonRooted;
+        } else {
+            // It is a "rooted" env.
+            if (locationFilters === undefined) {
+                return query.searchLocations === undefined;
             }
-            // `query.searchLocations` only had `null` in it.
-            return !query.searchLocations || query.searchLocations.length === 0;
+            // Check against the requested roots.
+            return locationFilters.some((filter) => filter(env.searchLocation!));
         }
-        // Check against the requested roots.
-        return locationFilters.some((filter) => filter(env.searchLocation!));
     }
     return (env) => {
         if (!checkKind(env)) {
@@ -65,6 +58,20 @@ export function getQueryFilter(query: PythonLocatorQuery): (env: PythonEnvInfo)
     };
 }
 
+function getSearchLocationFilters(query: PythonLocatorQuery): ((u: Uri) => boolean)[] | undefined {
+    if (query.searchLocations === undefined) {
+        return undefined;
+    }
+    const candidates: Uri[] = query.searchLocations.roots.filter((u) => !!u);
+    if (candidates.length === 0) {
+        return undefined;
+    }
+    return candidates.map((loc) => getURIFilter(loc, {
+        checkParent: true,
+        checkExact: true,
+    }));
+}
+
 /**
  * Unroll the given iterator into an array.
  *

diff --git a/src/client/pythonEnvironments/discovery/locators/index.ts b/src/client/pythonEnvironments/discovery/locators/index.ts
@@ -103,13 +103,12 @@ export class WorkspaceLocators extends Locator {
             return NOOP_ITERATOR;
         }
         const iterators = Object.keys(this.locators).map((key) => {
-            if (query?.searchLocations) {
+            if (query?.searchLocations !== undefined) {
                 const root = this.roots[key];
                 // Match any related search location.
                 const filter = getURIFilter(root, { checkParent: true, checkChild: true, checkExact: true });
                 // Ignore any requests for global envs.
-                const candidates: unknown = query.searchLocations.filter((u) => u !== null);
-                if (!(candidates as Uri[]).some(filter)) {
+                if (!query.searchLocations.roots.some(filter)) {
                     // This workspace folder did not match the query, so skip it!
                     return NOOP_ITERATOR;
                 }

diff --git a/src/client/pythonEnvironments/legacyIOC.ts b/src/client/pythonEnvironments/legacyIOC.ts
@@ -267,7 +267,7 @@ class ComponentAdapter implements IComponentAdapter {
         if (resource !== undefined) {
             const wsFolder = vscode.workspace.getWorkspaceFolder(resource);
             if (wsFolder !== undefined) {
-                query.searchLocations = [wsFolder.uri];
+                query.searchLocations = { roots: [wsFolder.uri] };
             }
         }
 

diff --git a/src/test/pythonEnvironments/base/locatorUtils.unit.test.ts b/src/test/pythonEnvironments/base/locatorUtils.unit.test.ts
@@ -73,8 +73,7 @@ suite('Python envs locator utils - getQueryFilter', () => {
         const queries: PythonLocatorQuery[] = [
             {},
             { kinds: [] },
-            { searchLocations: [] },
-            { kinds: [], searchLocations: [] },
+            // Any "defined" value for searchLocations causes filtering...
         ];
         queries.forEach((query) => {
             test(`all envs kept (query ${query})`, () => {
@@ -144,7 +143,11 @@ suite('Python envs locator utils - getQueryFilter', () => {
 
     suite('searchLocations', () => {
         test('match none', () => {
-            const query: PythonLocatorQuery = { searchLocations: [doesNotExist] };
+            const query: PythonLocatorQuery = {
+                searchLocations: {
+                    roots: [doesNotExist],
+                },
+            };
 
             const filter = getQueryFilter(query);
             const filtered = envs.filter(filter);
@@ -154,11 +157,13 @@ suite('Python envs locator utils - getQueryFilter', () => {
 
         test('match one (multiple locations)', () => {
             const expected = [envSL4];
-            const searchLocations = [
-                envSL4.searchLocation!,
-                doesNotExist,
-                envSL4.searchLocation!, // repeated
-            ];
+            const searchLocations = {
+                roots: [
+                    envSL4.searchLocation!,
+                    doesNotExist,
+                    envSL4.searchLocation!, // repeated
+                ],
+            };
             const query: PythonLocatorQuery = { searchLocations };
 
             const filter = getQueryFilter(query);
@@ -169,9 +174,11 @@ suite('Python envs locator utils - getQueryFilter', () => {
 
         test('match multiple (one location)', () => {
             const expected = [envS3, envSL2];
-            const searchLocations = [
-                workspaceRoot,
-            ];
+            const searchLocations = {
+                roots: [
+                    workspaceRoot,
+                ],
+            };
             const query: PythonLocatorQuery = { searchLocations };
 
             const filter = getQueryFilter(query);
@@ -182,8 +189,10 @@ suite('Python envs locator utils - getQueryFilter', () => {
 
         test('match multiple (multiple locations)', () => {
             const expected = [envS3, ...rootedLocatedEnvs];
-            const searchLocations = rootedLocatedEnvs.map((env) => env.searchLocation!);
-            searchLocations.push(doesNotExist);
+            const searchLocations = {
+                roots: rootedLocatedEnvs.map((env) => env.searchLocation!),
+            };
+            searchLocations.roots.push(doesNotExist);
             const query: PythonLocatorQuery = { searchLocations };
 
             const filter = getQueryFilter(query);
@@ -194,8 +203,11 @@ suite('Python envs locator utils - getQueryFilter', () => {
 
         test('match multiple (include non-searched envs)', () => {
             const expected = [...plainEnvs, ...locatedEnvs, envS3, ...rootedLocatedEnvs];
-            const searchLocations = [null, ...rootedLocatedEnvs.map((env) => env.searchLocation!)];
-            searchLocations.push(doesNotExist);
+            const searchLocations = {
+                roots: rootedLocatedEnvs.map((env) => env.searchLocation!),
+                includeNonRooted: true,
+            };
+            searchLocations.roots.push(doesNotExist);
             const query: PythonLocatorQuery = { searchLocations };
 
             const filter = getQueryFilter(query);
@@ -206,7 +218,9 @@ suite('Python envs locator utils - getQueryFilter', () => {
 
         test('match all searched', () => {
             const expected = [...rootedEnvs, ...rootedLocatedEnvs];
-            const searchLocations = expected.map((env) => env.searchLocation!);
+            const searchLocations = {
+                roots: expected.map((env) => env.searchLocation!),
+            };
             const query: PythonLocatorQuery = { searchLocations };
 
             const filter = getQueryFilter(query);
@@ -217,7 +231,10 @@ suite('Python envs locator utils - getQueryFilter', () => {
 
         test('match all (including non-searched)', () => {
             const expected = envs;
-            const searchLocations = [null, ...expected.map((env) => env.searchLocation!)];
+            const searchLocations = {
+                roots: expected.map((env) => env.searchLocation!),
+                includeNonRooted: true,
+            };
             const query: PythonLocatorQuery = { searchLocations };
 
             const filter = getQueryFilter(query);
@@ -228,7 +245,9 @@ suite('Python envs locator utils - getQueryFilter', () => {
 
         test('match all searched under one root', () => {
             const expected = [envS1, envS2, envSL1, envSL3, envSL5];
-            const searchLocations = [Uri.file(homeDir)];
+            const searchLocations = {
+                roots: [Uri.file(homeDir)],
+            };
             const query: PythonLocatorQuery = { searchLocations };
 
             const filter = getQueryFilter(query);
@@ -237,9 +256,12 @@ suite('Python envs locator utils - getQueryFilter', () => {
             assert.deepEqual(filtered, expected);
         });
 
-        test('match only non-searched envs (null location)', () => {
+        test('match only non-searched envs (empty roots)', () => {
             const expected = [...plainEnvs, ...locatedEnvs];
-            const searchLocations = [null];
+            const searchLocations = {
+                roots: [],
+                includeNonRooted: true,
+            };
             const query: PythonLocatorQuery = { searchLocations };
 
             const filter = getQueryFilter(query);
@@ -250,32 +272,26 @@ suite('Python envs locator utils - getQueryFilter', () => {
 
         test('match only non-searched envs (with unmatched location)', () => {
             const expected = [...plainEnvs, ...locatedEnvs];
-            const searchLocations: (Uri | null)[] = [null];
-            searchLocations.push(doesNotExist);
+            const searchLocations = {
+                roots: [doesNotExist],
+                includeNonRooted: true,
+            };
             const query: PythonLocatorQuery = { searchLocations };
 
             const filter = getQueryFilter(query);
             const filtered = envs.filter(filter);
 
             assert.deepEqual(filtered, expected);
         });
-
-        test('match only non-searched envs (explicit null)', () => {
-            const expected = [...plainEnvs, ...locatedEnvs];
-            const query: PythonLocatorQuery = { searchLocations: null };
-
-            const filter = getQueryFilter(query);
-            const filtered = envs.filter(filter);
-
-            assert.deepEqual(filtered, expected);
-        });
     });
 
     suite('mixed query', () => {
         test('match none', () => {
             const query: PythonLocatorQuery = {
                 kinds: [PythonEnvKind.OtherGlobal],
-                searchLocations: [doesNotExist],
+                searchLocations: {
+                    roots: [doesNotExist],
+                },
             };
 
             const filter = getQueryFilter(query);
@@ -287,8 +303,10 @@ suite('Python envs locator utils - getQueryFilter', () => {
         test('match some', () => {
             const expected = [envSL1, envSL4, envSL5];
             const kinds = [PythonEnvKind.Venv, PythonEnvKind.Custom];
-            const searchLocations = rootedLocatedEnvs.map((env) => env.searchLocation!);
-            searchLocations.push(doesNotExist);
+            const searchLocations = {
+                roots: rootedLocatedEnvs.map((env) => env.searchLocation!),
+            };
+            searchLocations.roots.push(doesNotExist);
             const query: PythonLocatorQuery = { kinds, searchLocations };
 
             const filter = getQueryFilter(query);
@@ -300,7 +318,9 @@ suite('Python envs locator utils - getQueryFilter', () => {
         test('match all', () => {
             const expected = [...rootedEnvs, ...rootedLocatedEnvs];
             const kinds: PythonEnvKind[] = getEnumValues(PythonEnvKind);
-            const searchLocations = expected.map((env) => env.searchLocation!);
+            const searchLocations = {
+                roots: expected.map((env) => env.searchLocation!),
+            };
             const query: PythonLocatorQuery = { kinds, searchLocations };
 
             const filter = getQueryFilter(query);

diff --git a/src/test/pythonEnvironments/base/locators.unit.test.ts b/src/test/pythonEnvironments/base/locators.unit.test.ts
@@ -96,7 +96,7 @@ suite('Python envs locators - Locators', () => {
         test('with query', async () => {
             const expected: PythonLocatorQuery = {
                 kinds: [PythonEnvKind.Venv],
-                searchLocations: [Uri.file('???')]
+                searchLocations: { roots: [Uri.file('???')] },
             };
             let query: PythonLocatorQuery | undefined;
             async function onQuery(q: PythonLocatorQuery | undefined, e: PythonEnvInfo[]) {

diff --git a/src/test/pythonEnvironments/discovery/locators/index.unit.test.ts b/src/test/pythonEnvironments/discovery/locators/index.unit.test.ts
@@ -392,8 +392,9 @@ suite('WorkspaceLocators', () => {
             ]);
             const folders = new WorkspaceFolders([root1, root2]);
             locators.activate(folders);
+            const query = { searchLocations: { roots: [root1] } };
 
-            const iterators = locators.iterEnvs({ searchLocations: [root1] });
+            const iterators = locators.iterEnvs(query);
             const envs = await getEnvs(iterators);
 
             expect(envs).to.deep.equal(expected);
@@ -417,8 +418,9 @@ suite('WorkspaceLocators', () => {
             ]);
             const folders = new WorkspaceFolders([root1, root2]);
             locators.activate(folders);
+            const query = { searchLocations: { roots: [root1, root2] } };
 
-            const iterators = locators.iterEnvs({ searchLocations: [root1, root2] });
+            const iterators = locators.iterEnvs(query);
             const envs = await getEnvs(iterators);
 
             expect(envs).to.deep.equal(expected);
@@ -441,8 +443,9 @@ suite('WorkspaceLocators', () => {
             ]);
             const folders = new WorkspaceFolders([root1, root2]);
             locators.activate(folders);
+            const query = { searchLocations: { roots: [Uri.file('baz')] } };
 
-            const iterators = locators.iterEnvs({ searchLocations: [Uri.file('baz')] });
+            const iterators = locators.iterEnvs(query);
             const envs = await getEnvs(iterators);
 
             expect(envs).to.deep.equal([]);