Rework search path resolution

src/Analysis/Ast/Impl/Modules/Resolution/MainModuleResolution.cs

@@ -37,7 +37,6 @@ namespace Microsoft.Python.Analysis.Modules.Resolution {
     internal sealed class MainModuleResolution : ModuleResolutionBase, IModuleManagement {
         private readonly ConcurrentDictionary<string, IPythonModule> _specialized = new ConcurrentDictionary<string, IPythonModule>();
         private IRunningDocumentTable _rdt;
-        private IReadOnlyList<string> _searchPaths;
 
         public MainModuleResolution(string root, IServiceContainer services)
             : base(root, services) { }
@@ -62,29 +61,6 @@ internal async Task InitializeAsync(CancellationToken cancellationToken = defaul
             cancellationToken.ThrowIfCancellationRequested();
         }
 
-        public async Task<IReadOnlyList<string>> GetSearchPathsAsync(CancellationToken cancellationToken = default) {
-            if (_searchPaths != null) {
-                return _searchPaths;
-            }
-
-            _searchPaths = await GetInterpreterSearchPathsAsync(cancellationToken);
-            Debug.Assert(_searchPaths != null, "Should have search paths");
-            _searchPaths = _searchPaths ?? Array.Empty<string>();
-
-            _log?.Log(TraceEventType.Verbose, "Python search paths:");
-            foreach (var s in _searchPaths) {
-                _log?.Log(TraceEventType.Verbose, $"    {s}");
-            }
-
-            var configurationSearchPaths = Configuration.SearchPaths ?? Array.Empty<string>();
-
-            _log?.Log(TraceEventType.Verbose, "Configuration search paths:");
-            foreach (var s in configurationSearchPaths) {
-                _log?.Log(TraceEventType.Verbose, $"    {s}");
-            }
-            return _searchPaths;
-        }
-
         protected override IPythonModule CreateModule(string name) {
             var moduleImport = CurrentPathResolver.GetModuleImportFromModuleName(name);
             if (moduleImport == null) {
@@ -135,11 +111,11 @@ protected override IPythonModule CreateModule(string name) {
             return GetRdt().AddModule(mco);
         }
 
-        private async Task<IReadOnlyList<string>> GetInterpreterSearchPathsAsync(CancellationToken cancellationToken = default) {
+        private async Task<IReadOnlyList<PythonLibraryPath>> GetInterpreterSearchPathsAsync(CancellationToken cancellationToken = default) {
             if (!_fs.FileExists(Configuration.InterpreterPath)) {
                 _log?.Log(TraceEventType.Warning, "Interpreter does not exist:", Configuration.InterpreterPath);
                 _ui?.ShowMessageAsync(Resources.InterpreterNotFound, TraceEventType.Error);
-                return Array.Empty<string>();
+                return Array.Empty<PythonLibraryPath>();
             }
 
             _log?.Log(TraceEventType.Information, "GetCurrentSearchPaths", Configuration.InterpreterPath);
@@ -148,11 +124,11 @@ private async Task<IReadOnlyList<string>> GetInterpreterSearchPathsAsync(Cancell
                 var ps = _services.GetService<IProcessServices>();
                 var paths = await PythonLibraryPath.GetSearchPathsAsync(Configuration, fs, ps, cancellationToken);
                 cancellationToken.ThrowIfCancellationRequested();
-                return paths.MaybeEnumerate().Select(p => p.Path).ToArray();
+                return paths.ToArray();
             } catch (InvalidOperationException ex) {
                 _log?.Log(TraceEventType.Warning, "Exception getting search paths", ex);
                 _ui?.ShowMessageAsync(Resources.ExceptionGettingSearchPaths, TraceEventType.Error);
-                return Array.Empty<string>();
+                return Array.Empty<PythonLibraryPath>();
             }
         }
 
@@ -211,16 +187,13 @@ public async Task ReloadAsync(CancellationToken cancellationToken = default) {
             var addedRoots = new HashSet<string>();
             addedRoots.UnionWith(PathResolver.SetRoot(Root));
 
-            InterpreterPaths = await GetSearchPathsAsync(cancellationToken);
+            var ps = _services.GetService<IProcessServices>();
 
-            IEnumerable<string> userSearchPaths = Configuration.SearchPaths;
-            InterpreterPaths = InterpreterPaths.Except(userSearchPaths, StringExtensions.PathsStringComparer).Where(p => !p.PathEquals(Root));
+            var paths = await GetInterpreterSearchPathsAsync(cancellationToken);
+            var (interpreterPaths, userPaths) = PythonLibraryPath.ClassifyPaths(Root, _fs, paths, Configuration.SearchPaths);
 
-            if (Root != null) {
-                var underRoot = userSearchPaths.ToLookup(p => _fs.IsPathUnderRoot(Root, p));
-                userSearchPaths = underRoot[true];
-                InterpreterPaths = underRoot[false].Concat(InterpreterPaths);
-            }
+            InterpreterPaths = interpreterPaths.Select(p => p.Path);
+            var userSearchPaths = userPaths.Select(p => p.Path);
 
             _log?.Log(TraceEventType.Information, "Interpreter search paths:");
             foreach (var s in InterpreterPaths) {

src/Analysis/Ast/Impl/get_search_paths.py

@@ -45,6 +45,11 @@ def clean(path):
 BEFORE_SITE = set(clean(p) for p in BEFORE_SITE)
 AFTER_SITE = set(clean(p) for p in AFTER_SITE)
 
+try:
+    SITE_PKGS = set(clean(p) for p in site.getsitepackages())
+except AttributeError:
+    SITE_PKGS = set()
+
 for prefix in [
     sys.prefix,
     sys.exec_prefix,
@@ -69,4 +74,7 @@ def clean(path):
         if p in BEFORE_SITE:
             print("%s|stdlib|" % p)
         elif p in AFTER_SITE:
-            print("%s||" % p)
+            if p in SITE_PKGS:
+                print("%s|site|" % p)
+            else:
+                print("%s|pth|" % p)

src/Analysis/Ast/Test/PathClassificationTests.cs

@@ -0,0 +1,218 @@
+// Copyright(c) Microsoft Corporation
+// All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the License); you may not use
+// this file except in compliance with the License. You may obtain a copy of the
+// License at http://www.apache.org/licenses/LICENSE-2.0
+//
+// THIS CODE IS PROVIDED ON AN  *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS
+// OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY
+// IMPLIED WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
+// MERCHANTABILITY OR NON-INFRINGEMENT.
+//
+// See the Apache Version 2.0 License for specific language governing
+// permissions and limitations under the License.
+
+using System;
+using System.IO;
+using FluentAssertions;
+using Microsoft.Python.Analysis.Core.Interpreter;
+using Microsoft.Python.Core.IO;
+using Microsoft.Python.Core.OS;
+using Microsoft.Python.Tests.Utilities.FluentAssertions;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using TestUtilities;
+
+namespace Microsoft.Python.Analysis.Tests {
+    [TestClass]
+    public class PathClassificationTests {
+        private readonly FileSystem _fs = new FileSystem(new OSPlatform());
+
+        public TestContext TestContext { get; set; }
+
+        [TestInitialize]
+        public void TestInitialize()
+            => TestEnvironmentImpl.TestInitialize($"{TestContext.FullyQualifiedTestClassName}.{TestContext.TestName}");
+
+        [TestCleanup]
+        public void Cleanup() => TestEnvironmentImpl.TestCleanup();
+
+        [TestMethod]
+        public void Plain() {
+            var appPath = TestData.GetTestSpecificPath("app.py");
+            var root = Path.GetDirectoryName(appPath);
+
+            var venv = Path.Combine(root, "venv");
+            var venvLib = Path.Combine(venv, "Lib");
+            var venvSitePackages = Path.Combine(venvLib, "site-packages");
+
+            var fromInterpreter = new[] {
+                new PythonLibraryPath(venvLib, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(venv, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(venvSitePackages, PythonLibraryPathType.Site),
+            };
+
+            var (interpreterPaths, userPaths) = PythonLibraryPath.ClassifyPaths(root, _fs, fromInterpreter, Array.Empty<string>());
+
+            interpreterPaths.Should().BeEquivalentToWithStrictOrdering(new[] {
+                new PythonLibraryPath(venvLib, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(venv, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(venvSitePackages, PythonLibraryPathType.Site),
+            });
+
+            userPaths.Should().BeEmpty();
+        }
+
+        [TestMethod]
+        public void WithSrcDir() {
+            var appPath = TestData.GetTestSpecificPath("app.py");
+            var root = Path.GetDirectoryName(appPath);
+
+            var venv = Path.Combine(root, "venv");
+            var venvLib = Path.Combine(venv, "Lib");
+            var venvSitePackages = Path.Combine(venvLib, "site-packages");
+
+            var src = Path.Combine(root, "src");
+
+            var fromInterpreter = new[] {
+                new PythonLibraryPath(venvLib, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(venv, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(venvSitePackages, PythonLibraryPathType.Site),
+            };
+
+            var fromUser = new[] {
+                "./src",
+            };
+
+            var (interpreterPaths, userPaths) = PythonLibraryPath.ClassifyPaths(root, _fs, fromInterpreter, fromUser);
+
+            interpreterPaths.Should().BeEquivalentToWithStrictOrdering(new[] {
+                new PythonLibraryPath(venvLib, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(venv, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(venvSitePackages, PythonLibraryPathType.Site),
+            });
+
+            userPaths.Should().BeEquivalentToWithStrictOrdering(new[] {
+                new PythonLibraryPath(src, PythonLibraryPathType.Unspecified),
+            });
+        }
+
+        [TestMethod]
+        public void NormalizeUser() {
+            var appPath = TestData.GetTestSpecificPath("app.py");
+            var root = Path.GetDirectoryName(appPath);
+
+            var src = Path.Combine(root, "src");
+
+            var fromUser = new[] {
+                "./src/",
+            };
+
+            var (interpreterPaths, userPaths) = PythonLibraryPath.ClassifyPaths(root, _fs, Array.Empty<PythonLibraryPath>(), fromUser);
+
+            interpreterPaths.Should().BeEmpty();
+
+            userPaths.Should().BeEquivalentToWithStrictOrdering(new[] {
+                new PythonLibraryPath(src, PythonLibraryPathType.Unspecified),
+            });
+        }
+
+        [TestMethod]
+        public void NestedUser() {
+            var appPath = TestData.GetTestSpecificPath("app.py");
+            var root = Path.GetDirectoryName(appPath);
+
+            var src = Path.Combine(root, "src");
+            var srcSomething = Path.Combine(src, "something");
+            var srcFoo = Path.Combine(src, "foo");
+            var srcFooBar = Path.Combine(srcFoo, "bar");
+
+            var fromUser = new[] {
+                "./src",
+                "./src/something",
+                "./src/foo/bar",
+                "./src/foo",
+            };
+
+            var (interpreterPaths, userPaths) = PythonLibraryPath.ClassifyPaths(root, _fs, Array.Empty<PythonLibraryPath>(), fromUser);
+
+            interpreterPaths.Should().BeEmpty();
+
+            userPaths.Should().BeEquivalentToWithStrictOrdering(new[] {
+                new PythonLibraryPath(src, PythonLibraryPathType.Unspecified),
+                new PythonLibraryPath(srcSomething, PythonLibraryPathType.Unspecified),
+                new PythonLibraryPath(srcFooBar, PythonLibraryPathType.Unspecified),
+                new PythonLibraryPath(srcFoo, PythonLibraryPathType.Unspecified),
+            });
+        }
+
+        [TestMethod]
+        public void NestedUserOrdering() {
+            var appPath = TestData.GetTestSpecificPath("app.py");
+            var root = Path.GetDirectoryName(appPath);
+
+            var src = Path.Combine(root, "src");
+            var srcSomething = Path.Combine(src, "something");
+            var srcFoo = Path.Combine(src, "foo");
+            var srcFooBar = Path.Combine(srcFoo, "bar");
+
+            var fromUser = new[] {
+                "./src/foo",
+                "./src/foo/bar",
+                "./src",
+                "./src/something",
+            };
+
+            var (interpreterPaths, userPaths) = PythonLibraryPath.ClassifyPaths(root, _fs, Array.Empty<PythonLibraryPath>(), fromUser);
+
+            interpreterPaths.Should().BeEmpty();
+
+            userPaths.Should().BeEquivalentToWithStrictOrdering(new[] {
+                new PythonLibraryPath(srcFoo, PythonLibraryPathType.Unspecified),
+                new PythonLibraryPath(srcFooBar, PythonLibraryPathType.Unspecified),
+                new PythonLibraryPath(src, PythonLibraryPathType.Unspecified),
+                new PythonLibraryPath(srcSomething, PythonLibraryPathType.Unspecified),
+            });
+        }
+
+        [TestMethod]
+        public void InsideStdLib() {
+            var appPath = TestData.GetTestSpecificPath("app.py");
+            var root = Path.GetDirectoryName(appPath);
+
+            var venv = Path.Combine(root, "venv");
+            var venvLib = Path.Combine(venv, "Lib");
+            var venvSitePackages = Path.Combine(venvLib, "site-packages");
+            var inside = Path.Combine(venvSitePackages, "inside");
+            var what = Path.Combine(venvSitePackages, "what");
+
+            var src = Path.Combine(root, "src");
+
+            var fromInterpreter = new[] {
+                new PythonLibraryPath(venvLib, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(venv, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(venvSitePackages, PythonLibraryPathType.Site),
+                new PythonLibraryPath(inside, PythonLibraryPathType.Pth),
+            };
+
+            var fromUser = new[] {
+                "./src",
+                "./venv/Lib/site-packages/what",
+            };
+
+            var (interpreterPaths, userPaths) = PythonLibraryPath.ClassifyPaths(root, _fs, fromInterpreter, fromUser);
+
+            interpreterPaths.Should().BeEquivalentToWithStrictOrdering(new[] {
+                new PythonLibraryPath(venvLib, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(venv, PythonLibraryPathType.StdLib),
+                new PythonLibraryPath(what, PythonLibraryPathType.Unspecified),
+                new PythonLibraryPath(venvSitePackages, PythonLibraryPathType.Site),
+                new PythonLibraryPath(inside, PythonLibraryPathType.Pth),
+            });
+
+            userPaths.Should().BeEquivalentToWithStrictOrdering(new[] {
+                new PythonLibraryPath(src, PythonLibraryPathType.Unspecified),
+            });
+        }
+    }
+}