Remove all references to dropped modules and force GC on reload

src/Analysis/Ast/Impl/Analyzer/Definitions/IPythonAnalyzer.cs

@@ -56,10 +56,11 @@ public interface IPythonAnalyzer {
         /// </summary>
         IReadOnlyList<DiagnosticsEntry> LintModule(IPythonModule module);
 
+
         /// <summary>
-        /// Removes all the modules from the analysis, except Typeshed and builtin
+        /// Removes all the modules from the analysis and restarts it, including stubs.
         /// </summary>
-        void ResetAnalyzer();
+        Task ResetAnalyzer();
 
         /// <summary>
         /// Returns list of currently loaded modules.

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

@@ -228,6 +228,8 @@ private void MergeStub() {
                 return;
             }
 
+            var builtins = Module.Interpreter.ModuleResolution.BuiltinsModule;
+
             // Note that scrape can pick up more functions than the stub contains
             // Or the stub can have definitions that scraping had missed. Therefore
             // merge is the combination of the two with the documentation coming
@@ -240,53 +242,101 @@ private void MergeStub() {
 
                 var sourceVar = Eval.GlobalScope.Variables[v.Name];
                 var sourceType = sourceVar?.Value.GetPythonType();
-                
+
                 // If stub says 'Any' but we have better type, keep the current type.
                 if (!IsStubBetterType(sourceType, stubType)) {
-                    continue;;
+                    continue;
                 }
 
-                // If types are the classes, merge members.
-                // Otherwise, replace type from one from the stub.
-                if (sourceType is PythonClassType cls && Module.Equals(cls.DeclaringModule)) {
-                    // If class exists and belongs to this module, add or replace
-                    // its members with ones from the stub, preserving documentation.
-                    // Don't augment types that do not come from this module.
-                    foreach (var name in stubType.GetMemberNames()) {
-                        var stubMember = stubType.GetMember(name);
-                        var member = cls.GetMember(name);
-
-                        var memberType = member?.GetPythonType();
-                        var stubMemberType = stubMember.GetPythonType();
-                        if (!IsStubBetterType(memberType, stubMemberType)) {
-                            continue;
+                // If type does not exist in module, but exists in stub, declare it unless it is an import.
+                // If types are the classes, merge members. Otherwise, replace type from one from the stub.
+                switch (sourceType) {
+                    case null:
+                        // Nothing in sources, but there is type in the stub. Declare it.
+                        if (v.Source == VariableSource.Declaration) {
+                            Eval.DeclareVariable(v.Name, v.Value, v.Source);
                         }
-
-                        // Get documentation from the current type, if any, since stubs
-                        // typically do not contain documentation while scraped code does.
-                        memberType?.TransferDocumentationAndLocation(stubMemberType);
-                        cls.AddMember(name, stubMember, overwrite: true);
-                    }
-                } else {
-                    // Re-declare variable with the data from the stub unless member is a module.
-                    // Modules members that are modules should remain as they are, i.e. os.path
-                    // should remain library with its own stub attached.
-                    if (!(stubType is IPythonModule)) {
-                        sourceType.TransferDocumentationAndLocation(stubType);
-                        // TODO: choose best type between the scrape and the stub. Stub probably should always win.
-                        var source = Eval.CurrentScope.Variables[v.Name]?.Source ?? VariableSource.Declaration;
-                        Eval.DeclareVariable(v.Name, v.Value, source);
-                    }
+                        break;
+
+                    case PythonClassType cls when Module.Equals(cls.DeclaringModule):
+                        // If class exists and belongs to this module, add or replace
+                        // its members with ones from the stub, preserving documentation.
+                        // Don't augment types that do not come from this module.
+                        // Do not replace __class__ since it has to match class type and we are not
+                        // replacing class type, we are only merging members.
+                        foreach (var name in stubType.GetMemberNames().Except(new[] { "__class__", "__base__", "__bases__", "__mro__", "mro" })) {
+                            var stubMember = stubType.GetMember(name);
+                            var member = cls.GetMember(name);
+
+                            var memberType = member?.GetPythonType();
+                            var stubMemberType = stubMember.GetPythonType();
+
+                            if (builtins.Equals(memberType?.DeclaringModule) || builtins.Equals(stubMemberType?.DeclaringModule)) {
+                                continue; // Leave builtins alone.
+                            }
+
+                            if (memberType?.DeclaringModule is SpecializedModule || stubMemberType?.DeclaringModule is SpecializedModule) {
+                                continue; // Leave specialized modules like typing alone.
+                            }
+
+                            if (!IsStubBetterType(memberType, stubMemberType)) {
+                                continue;
+                            }
+
+                            // Get documentation from the current type, if any, since stubs
+                            // typically do not contain documentation while scraped code does.
+                            TransferDocumentationAndLocation(memberType, stubMemberType);
+                            cls.AddMember(name, stubMember, overwrite: true);
+                        }
+                        break;
+
+                    case IPythonModule _:
+                        // We do not re-declare modules.
+                        break;
+
+                    default:
+                        // Re-declare variable with the data from the stub unless member is a module.
+                        // Modules members that are modules should remain as they are, i.e. os.path
+                        // should remain library with its own stub attached.
+                        var stubModule = stubType.DeclaringModule;
+                        if (!(stubType is IPythonModule) && !builtins.Equals(stubModule)) {
+                            TransferDocumentationAndLocation(sourceType, stubType);
+                            // TODO: choose best type between the scrape and the stub. Stub probably should always win.
+                            var source = Eval.CurrentScope.Variables[v.Name]?.Source ?? v.Source;
+                            Eval.DeclareVariable(v.Name, v.Value, source);
+                        }
+                        break;
                 }
             }
         }
 
         private static bool IsStubBetterType(IPythonType sourceType, IPythonType stubType) {
-            // If stub says 'Any' but we have better type, keep the current type.
             if (stubType.IsUnknown()) {
+                // Do not use worse types than what is in the module.
                 return false;
             }
-            return sourceType.IsUnknown() || !(stubType.DeclaringModule is TypingModule) || stubType.Name != "Any";
+            if (sourceType.IsUnknown()) {
+                return true; // Anything is better than unknowns.
+            }
+            // If stub says 'Any' but we have better type, keep the current type.
+            return !(stubType.DeclaringModule is TypingModule) || stubType.Name != "Any";
+        }
+
+        private static void TransferDocumentationAndLocation(IPythonType s, IPythonType d) {
+            if (s.IsUnknown()) {
+                return; // Do not transfer location of unknowns
+            }
+            // Documentation and location are always get transferred from module type
+            // to the stub type and never the other way around. This makes sure that
+            // we show documentation from the original module and goto definition
+            // navigates to the module source and not to the stub.
+            if (s != d && s is PythonType src && d is PythonType dst) {
+                var documentation = src.Documentation;
+                if (!string.IsNullOrEmpty(documentation)) {
+                    dst.SetDocumentation(documentation);
+                }
+                dst.Location = src.Location;
+            }
         }
     }
 }

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

@@ -17,7 +17,6 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
-using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Python.Analysis.Caching;
@@ -49,6 +48,7 @@ public sealed class PythonAnalyzer : IPythonAnalyzer, IDisposable {
         private int _version;
         private PythonAnalyzerSession _currentSession;
         private PythonAnalyzerSession _nextSession;
+        private bool _forceGCOnNextSession;
 
         public PythonAnalyzer(IServiceManager services, string cacheFolderPath = null) {
             _services = services;
@@ -196,16 +196,27 @@ public IReadOnlyList<DiagnosticsEntry> LintModule(IPythonModule module) {
             return optionsProvider?.Options?.LintingEnabled == false ? Array.Empty<DiagnosticsEntry>() : result;
         }
 
-        public void ResetAnalyzer() {
+        public async Task ResetAnalyzer() {
+            var interpreter = _services.GetService<IPythonInterpreter>();
+            var builtins = interpreter.ModuleResolution.BuiltinsModule;
+            builtins.SetAst(builtins.Analysis.Ast);
+
+            await interpreter.TypeshedResolution.ReloadAsync();
+            await interpreter.ModuleResolution.ReloadAsync();
+
             lock (_syncObj) {
-                _analysisEntries.Split(kvp => kvp.Key.IsTypeshed || kvp.Value.Module is IBuiltinsPythonModule, out var entriesToPreserve, out var entriesToRemove);
+                _forceGCOnNextSession = true;
+
+                _analysisEntries.Split(kvp => kvp.Value.Module is IBuiltinsPythonModule, out var entriesToPreserve, out var entriesToRemove);
                 _analysisEntries.Clear();
                 foreach (var (key, entry) in entriesToPreserve) {
                     _analysisEntries.Add(key, entry);
                 }
 
                 _dependencyResolver.RemoveKeys(entriesToRemove.Select(e => e.Key));
             }
+
+            _services.GetService<IRunningDocumentTable>().ReloadAll();
         }
 
         public IReadOnlyList<IPythonModule> LoadedModules {
@@ -218,7 +229,7 @@ public IReadOnlyList<IPythonModule> LoadedModules {
 
         public event EventHandler<AnalysisCompleteEventArgs> AnalysisComplete;
 
-        internal void RaiseAnalysisComplete(int moduleCount, double msElapsed) 
+        internal void RaiseAnalysisComplete(int moduleCount, double msElapsed)
             => AnalysisComplete?.Invoke(this, new AnalysisCompleteEventArgs(moduleCount, msElapsed));
 
         private void AnalyzeDocument(in AnalysisModuleKey key, in PythonAnalyzerEntry entry, in ImmutableArray<AnalysisModuleKey> dependencies) {
@@ -241,7 +252,7 @@ private void AnalyzeDocument(in AnalysisModuleKey key, in PythonAnalyzerEntry en
                 session.Start(true);
             }
         }
-        
+
         private bool TryCreateSession(in int graphVersion, in PythonAnalyzerEntry entry, out PythonAnalyzerSession session) {
             var analyzeUserModuleOutOfOrder = false;
             lock (_syncObj) {
@@ -276,7 +287,7 @@ private bool TryCreateSession(in int graphVersion, in PythonAnalyzerEntry entry,
                     session = null;
                     return false;
                 }
-                
+
                 if (_currentSession.IsCompleted) {
                     _currentSession = session = CreateSession(walker, null);
                     return true;
@@ -301,14 +312,21 @@ private void StartNextSession(Task task) {
             session.Start(false);
         }
 
-        private PythonAnalyzerSession CreateSession(in IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> walker, in PythonAnalyzerEntry entry) 
-            => new PythonAnalyzerSession(_services, _progress, _analysisCompleteEvent, _startNextSession, _disposeToken.CancellationToken, walker, _version, entry);
+        private PythonAnalyzerSession CreateSession(in IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> walker, in PythonAnalyzerEntry entry) {
+            bool forceGC;
+            lock (_syncObj) {
+                forceGC = _forceGCOnNextSession;
+                _forceGCOnNextSession = false;
+            }
+
+            return new PythonAnalyzerSession(_services, _progress, _analysisCompleteEvent, _startNextSession, _disposeToken.CancellationToken, walker, _version, entry, forceGC: forceGC);
+        }
 
         private void LoadMissingDocuments(IPythonInterpreter interpreter, ImmutableArray<AnalysisModuleKey> missingKeys) {
             if (missingKeys.Count == 0) {
                 return;
             }
-            
+
             var foundKeys = ImmutableArray<AnalysisModuleKey>.Empty;
             foreach (var missingKey in missingKeys) {
                 lock (_syncObj) {

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

@@ -45,7 +45,7 @@ internal sealed class PythonAnalyzerSession {
         private readonly IProgressReporter _progress;
         private readonly IPythonAnalyzer _analyzer;
         private readonly ILogger _log;
-        private readonly ITelemetryService _telemetry;
+        private readonly bool _forceGC;
 
         private State _state;
         private bool _isCanceled;
@@ -69,7 +69,8 @@ public PythonAnalyzerSession(IServiceManager services,
             CancellationToken analyzerCancellationToken,
             IDependencyChainWalker<AnalysisModuleKey, PythonAnalyzerEntry> walker,
             int version,
-            PythonAnalyzerEntry entry) {
+            PythonAnalyzerEntry entry,
+            bool forceGC = false) {
 
             _services = services;
             _analysisCompleteEvent = analysisCompleteEvent;
@@ -80,11 +81,11 @@ public PythonAnalyzerSession(IServiceManager services,
             _walker = walker;
             _entry = entry;
             _state = State.NotStarted;
+            _forceGC = forceGC;
 
             _diagnosticsService = _services.GetService<IDiagnosticsService>();
             _analyzer = _services.GetService<IPythonAnalyzer>();
             _log = _services.GetService<ILogger>();
-            _telemetry = _services.GetService<ITelemetryService>();
             _progress = progress;
         }
 
@@ -158,11 +159,12 @@ private async Task StartAsync() {
 
             var elapsed = stopWatch.Elapsed.TotalMilliseconds;
             LogResults(_log, elapsed, originalRemaining, remaining, Version);
-            ForceGCIfNeeded(originalRemaining, remaining);
+            ForceGCIfNeeded(_log, originalRemaining, remaining, _forceGC);
         }
 
-        private static void ForceGCIfNeeded(int originalRemaining, int remaining) {
-            if (originalRemaining - remaining > 1000) {
+        private static void ForceGCIfNeeded(ILogger logger, int originalRemaining, int remaining, bool force) {
+            if (force || originalRemaining - remaining > 1000) {
+                logger?.Log(TraceEventType.Verbose, "Forcing full garbage collection and heap compaction.");
                 GCSettings.LargeObjectHeapCompactionMode = GCLargeObjectHeapCompactionMode.CompactOnce;
                 GC.Collect();
             }
@@ -229,7 +231,7 @@ private async Task<int> AnalyzeAffectedEntriesAsync(Stopwatch stopWatch) {
         }
 
 
-        private bool IsAnalyzedLibraryInLoop(IDependencyChainNode<PythonAnalyzerEntry> node) 
+        private bool IsAnalyzedLibraryInLoop(IDependencyChainNode<PythonAnalyzerEntry> node)
             => !node.HasMissingDependencies && node.Value.IsAnalyzedLibrary(_walker.Version) && node.IsWalkedWithDependencies && node.IsValidVersion;
 
         private Task StartAnalysis(IDependencyChainNode<PythonAnalyzerEntry> node, AsyncCountdownEvent ace, Stopwatch stopWatch)
@@ -368,10 +370,22 @@ private void LogException(IPythonModule module, Exception exception) {
         }
 
         private IDocumentAnalysis CreateAnalysis(IDependencyChainNode<PythonAnalyzerEntry> node, IDocument document, PythonAst ast, int version, ModuleWalker walker, bool isCanceled) {
+            var canHaveLibraryAnalysis = false;
+
+            // Don't try to drop builtins; it causes issues elsewhere.
+            // We probably want the builtin module's AST and other info for evaluation.
+            switch (document.ModuleType) {
+                case ModuleType.Library:
+                case ModuleType.Stub:
+                case ModuleType.Compiled:
+                    canHaveLibraryAnalysis = true;
+                    break;
+            }
+
             var createLibraryAnalysis = !isCanceled &&
                 node != null &&
                 !node.HasMissingDependencies &&
-                document.ModuleType == ModuleType.Library &&
+                canHaveLibraryAnalysis &&
                 !document.IsOpen &&
                 node.HasOnlyWalkedDependencies &&
                 node.IsValidVersion;

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

@@ -13,10 +13,7 @@
 // See the Apache Version 2.0 License for specific language governing
 // permissions and limitations under the License.
 
-using System;
-
 namespace Microsoft.Python.Analysis.Modules {
-    [Flags]
     public enum ModuleType {
         /// <summary>
         /// Module is user file in the workspace.

src/Analysis/Ast/Impl/Types/PythonType.cs

@@ -17,6 +17,7 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
+using Microsoft.Python.Analysis.Modules;
 using Microsoft.Python.Analysis.Values;
 using Microsoft.Python.Core.Diagnostics;
 
@@ -51,6 +52,14 @@ private PythonType(string name, Location location, BuiltinTypeId typeId) : base(
 
         #region ILocatedMember
         public override PythonMemberType MemberType => _typeId.GetMemberId();
+
+        public override void AddReference(Location location) {
+            if (DeclaringModule == null || DeclaringModule.ModuleType == ModuleType.Builtins) {
+                return;
+            }
+
+            base.AddReference(location);
+        }
         #endregion
 
         #region IPythonType