Fix #1156: Unresolved import for builtin packages after reload (#1496)

* Fix #1156: Unresolved import for builtin packages after reload

* Address CR comments

Author: AlexanderSher

src/Analysis/Ast/Impl/Analyzer/PythonInterpreter.cs

@@ -33,6 +33,7 @@ namespace Microsoft.Python.Analysis.Analyzer {
     public sealed class PythonInterpreter : IPythonInterpreter {
         private MainModuleResolution _moduleResolution;
         private TypeshedResolution _stubResolution;
+        private IPythonType _unknownType;
         private readonly object _lock = new object();
 
         private readonly Dictionary<BuiltinTypeId, IPythonType> _builtinTypes = new Dictionary<BuiltinTypeId, IPythonType>();
@@ -42,30 +43,20 @@ private PythonInterpreter(InterpreterConfiguration configuration) {
             LanguageVersion = Configuration.Version.ToLanguageVersion();
         }
 
-        private async Task LoadBuiltinTypesAsync(string root, IServiceManager sm, CancellationToken cancellationToken = default) {
+        private async Task InitializeAsync(string root, IServiceManager sm, CancellationToken cancellationToken = default) {
             cancellationToken.ThrowIfCancellationRequested();
 
             sm.AddService(this);
             _moduleResolution = new MainModuleResolution(root, sm);
-            _stubResolution = new TypeshedResolution(sm);
-
-            await _moduleResolution.InitializeAsync(cancellationToken);
-            await _stubResolution.InitializeAsync(cancellationToken);
-
-            lock (_lock) {
-                var builtinModule = _moduleResolution.CreateBuiltinsModule();
-                _builtinTypes[BuiltinTypeId.NoneType]
-                    = new PythonType("NoneType", new Location(builtinModule), string.Empty, BuiltinTypeId.NoneType);
-                _builtinTypes[BuiltinTypeId.Unknown]
-                    = UnknownType = new PythonType("Unknown", new Location(builtinModule), string.Empty);
-            }
-
-            await _moduleResolution.LoadBuiltinTypesAsync(cancellationToken);
+            _stubResolution = new TypeshedResolution(Configuration.TypeshedPath, sm);
+            
+            await _stubResolution.ReloadAsync(cancellationToken);
+            await _moduleResolution.ReloadAsync(cancellationToken);
         }
 
         public static async Task<IPythonInterpreter> CreateAsync(InterpreterConfiguration configuration, string root, IServiceManager sm, CancellationToken cancellationToken = default) {
             var pi = new PythonInterpreter(configuration);
-            await pi.LoadBuiltinTypesAsync(root, sm, cancellationToken);
+            await pi.InitializeAsync(root, sm, cancellationToken);
 
             // Specialize typing
             TypingModule.Create(sm);
@@ -92,6 +83,24 @@ public static async Task<IPythonInterpreter> CreateAsync(InterpreterConfiguratio
         /// </summary>
         public IModuleResolution TypeshedResolution => _stubResolution;
 
+        /// <summary>
+        /// Unknown type.
+        /// </summary>
+        public IPythonType UnknownType {
+            get {
+                lock (_lock) {
+                    var type = _unknownType;
+                    if (type != null) {
+                        return type;
+                    }
+
+                    _unknownType = new PythonType("Unknown", new Location(_moduleResolution.BuiltinsModule), string.Empty);
+                    _builtinTypes[BuiltinTypeId.Unknown] = _unknownType;
+                    return _unknownType;
+                }
+            }
+        }
+
         /// <summary>
         /// Gets a well known built-in type such as int, list, dict, etc...
         /// </summary>
@@ -106,31 +115,34 @@ public IPythonType GetBuiltinType(BuiltinTypeId id) {
             }
 
             lock (_lock) {
-                if (_builtinTypes.TryGetValue(id, out var res) && res != null) {
-                    return res;
+                if (id == BuiltinTypeId.Unknown) {
+                    return UnknownType;
                 }
 
-                var bm = ModuleResolution.BuiltinsModule;
-                var typeName = id.GetTypeName(LanguageVersion);
-                if (typeName != null) {
-                    res = bm.GetMember(typeName) as IPythonType;
+                if (_builtinTypes.TryGetValue(id, out var type) && type != null) {
+                    return type;
                 }
 
-                if (res == null) {
-                    res = bm.GetAnyMember("__{0}__".FormatInvariant(id)) as IPythonType;
-                    if (res == null) {
-                        return _builtinTypes[BuiltinTypeId.Unknown];
+                if (id == BuiltinTypeId.NoneType) {
+                    type = new PythonType("NoneType", new Location(_moduleResolution.BuiltinsModule), string.Empty, BuiltinTypeId.NoneType);
+                } else {
+                    var bm = _moduleResolution.BuiltinsModule;
+                    var typeName = id.GetTypeName(LanguageVersion);
+                    if (typeName != null) {
+                        type = _moduleResolution.BuiltinsModule.GetMember(typeName) as IPythonType;
+                    }
+
+                    if (type == null) {
+                        type = bm.GetAnyMember("__{0}__".FormatInvariant(id)) as IPythonType;
+                        if (type == null) {
+                            return UnknownType;
+                        }
                     }
                 }
 
-                _builtinTypes[id] = res;
-                return res;
+                _builtinTypes[id] = type;
+                return type;
             }
         }
-
-        /// <summary>
-        /// Unknown type.
-        /// </summary>
-        public IPythonType UnknownType { get; private set; }
     }
 }

src/Analysis/Ast/Impl/Modules/Definitions/IModuleManagement.cs

@@ -26,6 +26,16 @@ namespace Microsoft.Python.Analysis.Modules {
     /// Represents module resolution and search subsystem.
     /// </summary>
     public interface IModuleManagement: IModuleResolution {
+        /// <summary>
+        /// Builtins module name.
+        /// </summary>
+        string BuiltinModuleName { get; }
+
+        /// <summary>
+        /// Builtins module.
+        /// </summary>
+        IBuiltinsPythonModule BuiltinsModule { get; }
+
         /// <summary>
         /// Locates module by path.
         /// </summary>

src/Analysis/Ast/Impl/Modules/Definitions/IModuleResolution.cs

@@ -23,16 +23,6 @@ namespace Microsoft.Python.Analysis.Modules {
     /// Represents basic module resolution and search subsystem.
     /// </summary>
     public interface IModuleResolution {
-        /// <summary>
-        /// Builtins module name.
-        /// </summary>
-        string BuiltinModuleName { get; }
-
-        /// <summary>
-        /// Builtins module.
-        /// </summary>
-        IBuiltinsPythonModule BuiltinsModule { get; }
-
         /// <summary>
         /// Path resolver providing file resolution in module imports.
         /// </summary>

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

@@ -33,41 +33,30 @@
 using Microsoft.Python.Core.Diagnostics;
 using Microsoft.Python.Core.IO;
 using Microsoft.Python.Core.OS;
+using Microsoft.Python.Core.Services;
 
 namespace Microsoft.Python.Analysis.Modules.Resolution {
     internal sealed class MainModuleResolution : ModuleResolutionBase, IModuleManagement {
         private readonly ConcurrentDictionary<string, IPythonModule> _specialized = new ConcurrentDictionary<string, IPythonModule>();
+        private readonly IUIService _ui;
         private IRunningDocumentTable _rdt;
 
         private ImmutableArray<string> _userPaths = ImmutableArray<string>.Empty;
 
         public MainModuleResolution(string root, IServiceContainer services)
-            : base(root, services) { }
+            : base(root, services) {
 
-        internal IBuiltinsPythonModule CreateBuiltinsModule() {
-            if (BuiltinsModule == null) {
-                // Initialize built-in
-                var moduleName = BuiltinTypeId.Unknown.GetModuleName(_interpreter.LanguageVersion);
-
-                StubCache = _services.GetService<IStubCache>();
-                var modulePath = StubCache.GetCacheFilePath(_interpreter.Configuration.InterpreterPath);
-
-                var b = new BuiltinsPythonModule(moduleName, modulePath, _services);
-                BuiltinsModule = b;
-                Modules[BuiltinModuleName] = new ModuleRef(b);
-            }
-            return BuiltinsModule;
+            _ui = services.GetService<IUIService>();
         }
+        
+        public string BuiltinModuleName => BuiltinTypeId.Unknown.GetModuleName(Interpreter.LanguageVersion);
 
-        internal async Task InitializeAsync(CancellationToken cancellationToken = default) {
-            await ReloadAsync(cancellationToken);
-            cancellationToken.ThrowIfCancellationRequested();
-        }
+        public IBuiltinsPythonModule BuiltinsModule { get; private set; }
 
         protected override IPythonModule CreateModule(string name) {
             var moduleImport = CurrentPathResolver.GetModuleImportFromModuleName(name);
             if (moduleImport == null) {
-                _log?.Log(TraceEventType.Verbose, "Import not found: ", name);
+                Log?.Log(TraceEventType.Verbose, "Import not found: ", name);
                 return null;
             }
 
@@ -83,7 +72,7 @@ protected override IPythonModule CreateModule(string name) {
             // First check stub next to the module.
             if (!TryCreateModuleStub(name, moduleImport.ModulePath, out var stub)) {
                 // If nothing found, try Typeshed.
-                stub = _interpreter.TypeshedResolution.GetOrLoadModule(moduleImport.IsBuiltin ? name : moduleImport.FullName);
+                stub = Interpreter.TypeshedResolution.GetOrLoadModule(moduleImport.IsBuiltin ? name : moduleImport.FullName);
             }
 
             // If stub is created and its path equals to module, return that stub as module
@@ -92,16 +81,16 @@ protected override IPythonModule CreateModule(string name) {
             }
 
             if (moduleImport.IsBuiltin) {
-                _log?.Log(TraceEventType.Verbose, "Create built-in compiled (scraped) module: ", name, Configuration.InterpreterPath);
-                return new CompiledBuiltinPythonModule(name, stub, _services);
+                Log?.Log(TraceEventType.Verbose, "Create built-in compiled (scraped) module: ", name, Configuration.InterpreterPath);
+                return new CompiledBuiltinPythonModule(name, stub, Services);
             }
 
             if (moduleImport.IsCompiled) {
-                _log?.Log(TraceEventType.Verbose, "Create compiled (scraped): ", moduleImport.FullName, moduleImport.ModulePath, moduleImport.RootPath);
-                return new CompiledPythonModule(moduleImport.FullName, ModuleType.Compiled, moduleImport.ModulePath, stub, _services);
+                Log?.Log(TraceEventType.Verbose, "Create compiled (scraped): ", moduleImport.FullName, moduleImport.ModulePath, moduleImport.RootPath);
+                return new CompiledPythonModule(moduleImport.FullName, ModuleType.Compiled, moduleImport.ModulePath, stub, Services);
             }
 
-            _log?.Log(TraceEventType.Verbose, "Import: ", moduleImport.FullName, moduleImport.ModulePath);
+            Log?.Log(TraceEventType.Verbose, "Import: ", moduleImport.FullName, moduleImport.ModulePath);
             // Module inside workspace == user code.
 
             var mco = new ModuleCreationOptions {
@@ -114,24 +103,22 @@ protected override IPythonModule CreateModule(string name) {
             return GetRdt().AddModule(mco);
         }
 
-        private async Task<IReadOnlyList<PythonLibraryPath>> GetInterpreterSearchPathsAsync(CancellationToken cancellationToken = default) {
-            if (!_fs.FileExists(Configuration.InterpreterPath)) {
-                _log?.Log(TraceEventType.Warning, "Interpreter does not exist:", Configuration.InterpreterPath);
+        private async Task<ImmutableArray<PythonLibraryPath>> GetInterpreterSearchPathsAsync(CancellationToken cancellationToken = default) {
+            if (!FileSystem.FileExists(Configuration.InterpreterPath)) {
+                Log?.Log(TraceEventType.Warning, "Interpreter does not exist:", Configuration.InterpreterPath);
                 _ui?.ShowMessageAsync(Resources.InterpreterNotFound, TraceEventType.Error);
-                return Array.Empty<PythonLibraryPath>();
+                return ImmutableArray<PythonLibraryPath>.Empty;
             }
 
-            _log?.Log(TraceEventType.Information, "GetCurrentSearchPaths", Configuration.InterpreterPath);
+            Log?.Log(TraceEventType.Information, "GetCurrentSearchPaths", Configuration.InterpreterPath);
             try {
-                var fs = _services.GetService<IFileSystem>();
-                var ps = _services.GetService<IProcessServices>();
-                var paths = await PythonLibraryPath.GetSearchPathsAsync(Configuration, fs, ps, cancellationToken);
-                cancellationToken.ThrowIfCancellationRequested();
-                return paths.ToArray();
+                var fs = Services.GetService<IFileSystem>();
+                var ps = Services.GetService<IProcessServices>();
+                return await PythonLibraryPath.GetSearchPathsAsync(Configuration, fs, ps, cancellationToken);
             } catch (InvalidOperationException ex) {
-                _log?.Log(TraceEventType.Warning, "Exception getting search paths", ex);
+                Log?.Log(TraceEventType.Warning, "Exception getting search paths", ex);
                 _ui?.ShowMessageAsync(Resources.ExceptionGettingSearchPaths, TraceEventType.Error);
-                return Array.Empty<PythonLibraryPath>();
+                return ImmutableArray<PythonLibraryPath>.Empty;
             }
         }
 
@@ -157,8 +144,8 @@ public IPythonModule SpecializeModule(string name, Func<string, IPythonModule> s
         public IPythonModule GetSpecializedModule(string name)
         => _specialized.TryGetValue(name, out var module) ? module : null;
 
-        internal async Task LoadBuiltinTypesAsync(CancellationToken cancellationToken = default) {
-            var analyzer = _services.GetService<IPythonAnalyzer>();
+        internal async Task AddBuiltinTypesToPathResolverAsync(CancellationToken cancellationToken = default) {
+            var analyzer = Services.GetService<IPythonAnalyzer>();
             await analyzer.GetAnalysisAsync(BuiltinsModule, -1, cancellationToken);
 
             Check.InvalidOperation(!(BuiltinsModule.Analysis is EmptyAnalysis), "After await");
@@ -172,26 +159,6 @@ internal async Task LoadBuiltinTypesAsync(CancellationToken cancellationToken =
             }
         }
 
-        internal async Task ReloadSearchPaths(CancellationToken cancellationToken = default) {
-            var ps = _services.GetService<IProcessServices>();
-
-            var paths = await GetInterpreterSearchPathsAsync(cancellationToken);
-            var (interpreterPaths, userPaths) = PythonLibraryPath.ClassifyPaths(Root, _fs, paths, Configuration.SearchPaths);
-
-            InterpreterPaths = interpreterPaths.Select(p => p.Path);
-            _userPaths = userPaths.Select(p => p.Path);
-
-            _log?.Log(TraceEventType.Information, "Interpreter search paths:");
-            foreach (var s in InterpreterPaths) {
-                _log?.Log(TraceEventType.Information, $"    {s}");
-            }
-
-            _log?.Log(TraceEventType.Information, "User search paths:");
-            foreach (var s in _userPaths) {
-                _log?.Log(TraceEventType.Information, $"    {s}");
-            }
-        }
-
         public async Task ReloadAsync(CancellationToken cancellationToken = default) {
             foreach (var uri in Modules
                 .Where(m => m.Value.Value?.Name != BuiltinModuleName)
@@ -201,45 +168,78 @@ public async Task ReloadAsync(CancellationToken cancellationToken = default) {
             }
 
             // Preserve builtins, they don't need to be reloaded since interpreter does not change.
-            if (Modules.TryGetValue(BuiltinModuleName, out var builtins)) {
-                Modules.Clear();
-                Modules[BuiltinModuleName] = builtins;
-            }
+            var builtinsIsCreated = Modules.TryGetValue(BuiltinModuleName, out var builtinsRef);
+            Modules.Clear();
 
             await ReloadSearchPaths(cancellationToken);
 
-            PathResolver = new PathResolver(_interpreter.LanguageVersion, Root, InterpreterPaths, _userPaths);
+            PathResolver = new PathResolver(Interpreter.LanguageVersion, Root, InterpreterPaths, _userPaths);
 
             var addedRoots = new HashSet<string> { Root };
             addedRoots.UnionWith(InterpreterPaths);
             addedRoots.UnionWith(_userPaths);
             ReloadModulePaths(addedRoots);
+
+            if (!builtinsIsCreated) {
+                var builtinsModule = CreateBuiltinsModule(Services, Interpreter, StubCache);
+                BuiltinsModule = builtinsModule;
+                builtinsRef = new ModuleRef(builtinsModule);
+            }
+
+            Modules[BuiltinModuleName] = builtinsRef;
+            await AddBuiltinTypesToPathResolverAsync(cancellationToken);
+        }
+
+        private static IBuiltinsPythonModule CreateBuiltinsModule(IServiceContainer services, IPythonInterpreter interpreter, IStubCache stubCache) {
+            var moduleName = BuiltinTypeId.Unknown.GetModuleName(interpreter.LanguageVersion);
+            var modulePath = stubCache.GetCacheFilePath(interpreter.Configuration.InterpreterPath);
+            return new BuiltinsPythonModule(moduleName, modulePath, services);
+        }
+
+        private async Task ReloadSearchPaths(CancellationToken cancellationToken = default) {
+            var paths = await GetInterpreterSearchPathsAsync(cancellationToken);
+            var (interpreterPaths, userPaths) = PythonLibraryPath.ClassifyPaths(Root, FileSystem, paths, Configuration.SearchPaths);
+
+            InterpreterPaths = interpreterPaths.Select(p => p.Path);
+            _userPaths = userPaths.Select(p => p.Path);
+
+            if (Log != null) {
+                Log.Log(TraceEventType.Information, "Interpreter search paths:");
+                foreach (var s in InterpreterPaths) {
+                    Log.Log(TraceEventType.Information, $"    {s}");
+                }
+
+                Log.Log(TraceEventType.Information, "User search paths:");
+                foreach (var s in _userPaths) {
+                    Log.Log(TraceEventType.Information, $"    {s}");
+                }
+            }
         }
 
         public bool TryAddModulePath(in string path, in long fileSize, in bool allowNonRooted, out string fullModuleName)
             => PathResolver.TryAddModulePath(path, fileSize, allowNonRooted, out fullModuleName);
 
         // For tests
         internal void AddUnimportableModule(string moduleName)
-            => Modules[moduleName] = new ModuleRef(new SentinelModule(moduleName, _services));
+            => Modules[moduleName] = new ModuleRef(new SentinelModule(moduleName, Services));
 
         private bool TryCreateModuleStub(string name, string modulePath, out IPythonModule module) {
             // First check stub next to the module.
             if (!string.IsNullOrEmpty(modulePath)) {
                 var pyiPath = Path.ChangeExtension(modulePath, "pyi");
-                if (_fs.FileExists(pyiPath)) {
-                    module = new StubPythonModule(name, pyiPath, false, _services);
+                if (FileSystem.FileExists(pyiPath)) {
+                    module = new StubPythonModule(name, pyiPath, false, Services);
                     return true;
                 }
             }
 
             // Try location of stubs that are in a separate folder next to the package.
-            var stubPath = CurrentPathResolver.GetPossibleModuleStubPaths(name).FirstOrDefault(p => _fs.FileExists(p));
-            module = !string.IsNullOrEmpty(stubPath) ? new StubPythonModule(name, stubPath, false, _services) : null;
+            var stubPath = CurrentPathResolver.GetPossibleModuleStubPaths(name).FirstOrDefault(p => FileSystem.FileExists(p));
+            module = !string.IsNullOrEmpty(stubPath) ? new StubPythonModule(name, stubPath, false, Services) : null;
             return module != null;
         }
 
         private IRunningDocumentTable GetRdt()
-            => _rdt ?? (_rdt = _services.GetService<IRunningDocumentTable>());
+            => _rdt ?? (_rdt = Services.GetService<IRunningDocumentTable>());
     }
 }