Fix #825: Unresolved import not going away automatically after installing packages

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

@@ -54,5 +54,10 @@ public interface IPythonAnalyzer {
         /// Runs linters on the modules
         /// </summary>
         IReadOnlyList<DiagnosticsEntry> LintModule(IPythonModule module);
+
+        /// <summary>
+        /// Removes all the modules from the analysis, except Typeshed and builtin
+        /// </summary>
+        void ResetAnalyzer();
     }
 }

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

@@ -177,6 +177,17 @@ public IReadOnlyList<DiagnosticsEntry> LintModule(IPythonModule module) {
             return new LinterAggregator().Lint(module.Analysis, _services);
         }
 
+        public void ResetAnalyzer() {
+            lock (_syncObj) {
+                _analysisEntries.Split(kvp => kvp.Key.IsTypeshed || 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));
+            }
+        }
 
         private void AnalyzeDocument(AnalysisModuleKey key, PythonAnalyzerEntry entry, ImmutableArray<AnalysisModuleKey> dependencies, CancellationToken cancellationToken) {
             _analysisCompleteEvent.Reset();

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

@@ -140,7 +140,7 @@ private async Task StartAsync(IDependencyChainWalker<AnalysisModuleKey, PythonAn
                 }
 
                 if (!isCanceled) {
-                    _progress.ReportRemaining(walker.Remaining);
+                    _progress.ReportRemaining(remaining);
                 }
 
                 if (_log != null) {

src/Analysis/Ast/Impl/Dependencies/DependencyGraph.cs

@@ -15,6 +15,7 @@
 
 using System.Collections.Generic;
 using System.Linq;
+using Microsoft.Python.Core;
 using Microsoft.Python.Core.Collections;
 
 namespace Microsoft.Python.Analysis.Dependencies {
@@ -64,7 +65,7 @@ public DependencyVertex<TKey, TValue> AddOrUpdate(TKey key, TValue value, Immuta
 
             _verticesByKey[key] = changedVertex;
 
-            var missingKeysHashSet = new HashSet<TKey>();
+
             var vertices = _verticesByIndex
                 .Where(v => !v.IsSealed || v.HasMissingKeys)
                 .Select(v => GetOrCreateNonSealedVertex(version, v.Index))
@@ -75,6 +76,33 @@ public DependencyVertex<TKey, TValue> AddOrUpdate(TKey key, TValue value, Immuta
                 return changedVertex;
             }
 
+            CreateNewSnapshot(vertices, version);
+
+            return changedVertex;
+        }
+
+        public void RemoveKeys(ImmutableArray<TKey> keys) {
+            var version = Snapshot.Version + 1;
+
+            _verticesByIndex.Clear();
+            foreach (var key in keys) {
+                _verticesByKey.Remove(key);
+            }
+
+            foreach (var (key, currentVertex) in _verticesByKey) {
+                var changedVertex = new DependencyVertex<TKey, TValue>(key, currentVertex.Value, currentVertex.IncomingKeys, version, _verticesByIndex.Count);
+                _verticesByIndex.Add(changedVertex);
+            }
+
+            foreach (var vertex in _verticesByIndex) {
+                _verticesByKey[vertex.Key] = vertex;
+            }
+
+            CreateNewSnapshot(_verticesByIndex, version);
+        }
+
+        private void CreateNewSnapshot(IEnumerable<DependencyVertex<TKey, TValue>> vertices, int version) {
+            var missingKeysHashSet = new HashSet<TKey>();
             foreach (var vertex in vertices) {
                 var newIncoming = ImmutableArray<int>.Empty;
                 var oldIncoming = vertex.Incoming;
@@ -105,11 +133,9 @@ public DependencyVertex<TKey, TValue> AddOrUpdate(TKey key, TValue value, Immuta
                 vertex.Seal();
             }
 
-            Snapshot = new DependencyGraphSnapshot<TKey, TValue>(version, 
+            Snapshot = new DependencyGraphSnapshot<TKey, TValue>(version,
                 ImmutableArray<DependencyVertex<TKey, TValue>>.Create(_verticesByIndex),
                 ImmutableArray<TKey>.Create(missingKeysHashSet));
-
-            return changedVertex;
         }
 
         private DependencyVertex<TKey, TValue> GetOrCreateNonSealedVertex(int version, int index) {

src/Analysis/Ast/Impl/Dependencies/DependencyResolver.cs

@@ -29,18 +29,49 @@ public IDependencyChainWalker<TKey, TValue> NotifyChanges(TKey key, TValue value
             => NotifyChanges(key, value, ImmutableArray<TKey>.Create(incomingKeys));
 
         public IDependencyChainWalker<TKey, TValue> NotifyChanges(TKey key, TValue value, ImmutableArray<TKey> incomingKeys) {
-            int version;
-            ImmutableArray<DependencyVertex<TKey, TValue>> changedVertices;
-            DependencyGraphSnapshot<TKey, TValue> snapshot;
+            AddOrUpdate(key, value, incomingKeys, out var snapshot, out var changedVertices);
+            return CreateDependencyChainWalker(this, snapshot, changedVertices);
+        }
+
+        public IDependencyChainWalker<TKey, TValue> RemoveKeys(params TKey[] keys)
+            => RemoveKeys(ImmutableArray<TKey>.Create(keys));
 
+        public IDependencyChainWalker<TKey, TValue> RemoveKeys(ImmutableArray<TKey> keys) {
+            Remove(keys, out var snapshot, out var changedVertices);
+            return CreateDependencyChainWalker(this, snapshot, changedVertices);
+        }
+
+        private void AddOrUpdate(TKey key, TValue value, ImmutableArray<TKey> incomingKeys, out DependencyGraphSnapshot<TKey, TValue> snapshot, out ImmutableArray<DependencyVertex<TKey, TValue>> changedVertices) {
             lock (_syncObj) {
                 var dependencyVertex = _vertices.AddOrUpdate(key, value, incomingKeys);
                 snapshot = _vertices.Snapshot;
 
-                version = _vertices.Snapshot.Version;
                 _changedVertices[key] = dependencyVertex;
                 changedVertices = ImmutableArray<DependencyVertex<TKey, TValue>>.Create(_changedVertices.Values);
             }
+        }
+
+        private void Remove(ImmutableArray<TKey> keys, out DependencyGraphSnapshot<TKey, TValue> snapshot, out ImmutableArray<DependencyVertex<TKey, TValue>> changedVertices) {
+            lock (_syncObj) {
+                _vertices.RemoveKeys(keys);
+                snapshot = _vertices.Snapshot;
+
+                foreach (var key in keys) {
+                    _changedVertices.Remove(key);
+                }
+
+                foreach (var vertex in snapshot.Vertices.Where(v => v.HasMissingKeys)) {
+                    _changedVertices[vertex.Key] = vertex;
+                }
+
+                changedVertices = ImmutableArray<DependencyVertex<TKey, TValue>>.Create(_changedVertices.Values);
+            }
+        }
+
+        private static IDependencyChainWalker<TKey, TValue> CreateDependencyChainWalker(
+            DependencyResolver<TKey, TValue> dependencyResolver,
+            DependencyGraphSnapshot<TKey, TValue> snapshot,
+            ImmutableArray<DependencyVertex<TKey, TValue>> changedVertices) {
 
             var walkingGraph = CreateWalkingGraph(snapshot.Vertices, changedVertices);
             var affectedValues = walkingGraph.Select(v => v.DependencyVertex.Value);
@@ -52,7 +83,7 @@ public IDependencyChainWalker<TKey, TValue> NotifyChanges(TKey key, TValue value
                 vertex.SecondPass?.Seal();
             }
 
-            return new DependencyChainWalker(this, startingVertices, affectedValues, snapshot.MissingKeys, totalNodesCount, version);
+            return new DependencyChainWalker(dependencyResolver, startingVertices, affectedValues, snapshot.MissingKeys, totalNodesCount, snapshot.Version);
         }
 
         private void CommitChanges(DependencyVertex<TKey, TValue> vertex) {
@@ -63,7 +94,7 @@ private void CommitChanges(DependencyVertex<TKey, TValue> vertex) {
             }
         }
 
-        private ImmutableArray<WalkingVertex<TKey, TValue>> CreateWalkingGraph(ImmutableArray<DependencyVertex<TKey, TValue>> snapshot, ImmutableArray<DependencyVertex<TKey, TValue>> changedVertices) {
+        private static ImmutableArray<WalkingVertex<TKey, TValue>> CreateWalkingGraph(ImmutableArray<DependencyVertex<TKey, TValue>> snapshot, ImmutableArray<DependencyVertex<TKey, TValue>> changedVertices) {
             var analysisGraph = ImmutableArray<WalkingVertex<TKey, TValue>>.Empty;
             var nodesByVertexIndex = new Dictionary<int, WalkingVertex<TKey, TValue>>();
 
@@ -92,7 +123,7 @@ private ImmutableArray<WalkingVertex<TKey, TValue>> CreateWalkingGraph(Immutable
             return analysisGraph;
         }
 
-        private int FindLoops(ImmutableArray<WalkingVertex<TKey, TValue>> graph) {
+        private static int FindLoops(ImmutableArray<WalkingVertex<TKey, TValue>> graph) {
             var index = 0;
             var loopNumber = 0;
             var stackP = new Stack<WalkingVertex<TKey, TValue>>();
@@ -107,7 +138,7 @@ private int FindLoops(ImmutableArray<WalkingVertex<TKey, TValue>> graph) {
             return loopNumber;
         }
 
-        private void CheckForLoop(WalkingVertex<TKey, TValue> vertex, Stack<WalkingVertex<TKey, TValue>> stackP, Stack<WalkingVertex<TKey, TValue>> stackS, ref int counter, ref int loopNumber) {
+        private static void CheckForLoop(WalkingVertex<TKey, TValue> vertex, Stack<WalkingVertex<TKey, TValue>> stackP, Stack<WalkingVertex<TKey, TValue>> stackS, ref int counter, ref int loopNumber) {
             vertex.Index = counter++;
             stackP.Push(vertex);
             stackS.Push(vertex);

src/Analysis/Ast/Impl/Dependencies/IDependencyResolver.cs

@@ -27,5 +27,6 @@ namespace Microsoft.Python.Analysis.Dependencies {
     /// </summary>
     internal interface IDependencyResolver<TKey, TValue> {
         IDependencyChainWalker<TKey, TValue> NotifyChanges(TKey key, TValue value, ImmutableArray<TKey> incomingKeys);
+        IDependencyChainWalker<TKey, TValue> RemoveKeys(ImmutableArray<TKey> keys);
     }
 }

src/Analysis/Ast/Impl/Documents/Definitions/IDocument.cs

@@ -68,7 +68,7 @@ public interface IDocument: IPythonModule, IDisposable {
         void Update(IEnumerable<DocumentChange> changes);
 
         /// <summary>
-        /// Resets document buffer to the provided content.
+        /// Resets document buffer to the provided content or tries to load it if content is null, then parses and analyzes document.
         /// </summary>
         void Reset(string content);
 

src/Analysis/Ast/Impl/Documents/DocumentBuffer.cs

@@ -29,8 +29,8 @@ internal sealed class DocumentBuffer {
 
         public void Reset(int version, string content) {
             Version = version;
+            _content = content ?? string.Empty;
             _sb = null;
-            _content = content;
         }
 
         public void Update(IEnumerable<DocumentChange> changes) {

src/Analysis/Ast/Impl/Documents/RunningDocumentTable.cs

@@ -225,6 +225,8 @@ private DocumentEntry CreateDocument(ModuleCreationOptions mco) {
             var entry = new DocumentEntry(document);
             _documentsByUri[document.Uri] = entry;
             _documentsByName[mco.ModuleName] = entry;
+
+            _services.GetService<IPythonAnalyzer>().InvalidateAnalysis(document);
             return entry;
         }
 

src/Analysis/Ast/Impl/Modules/PythonModule.cs

@@ -111,7 +111,7 @@ internal PythonModule(ModuleCreationOptions creationOptions, IServiceContainer s
             if (ModuleType == ModuleType.Specialized || ModuleType == ModuleType.Unresolved) {
                 ContentState = State.Analyzed;
             }
-            InitializeContent(creationOptions.Content);
+            InitializeContent(creationOptions.Content, 0);
         }
 
         #region IPythonType
@@ -184,18 +184,6 @@ public virtual IEnumerable<string> GetMemberNames() {
         /// </summary>
         public IGlobalScope GlobalScope { get; private set; }
 
-        /// <summary>
-        /// Ensures that module content is loaded and analysis has started.
-        /// Typically module content is loaded at the creation time, but delay
-        /// loaded (lazy) modules may choose to defer content retrieval and
-        /// analysis until later time, when module members are actually needed.
-        /// </summary>
-        public async Task LoadAndAnalyzeAsync(CancellationToken cancellationToken = default) {
-            InitializeContent(null);
-            await GetAstAsync(cancellationToken);
-            await Services.GetService<IPythonAnalyzer>().GetAnalysisAsync(this, -1, cancellationToken);
-        }
-
         /// <summary>
         /// If module is a stub points to the primary module.
         /// Typically used in code navigation scenarios when user
@@ -215,27 +203,22 @@ protected virtual string LoadContent() {
             return null; // Keep content as null so module can be loaded later.
         }
 
-        private void InitializeContent(string content) {
-            bool startParse;
+        private void InitializeContent(string content, int version) {
             lock (AnalysisLock) {
-                LoadContent(content);
+                LoadContent(content, version);
 
-                startParse = ContentState < State.Parsing && _parsingTask == null;
+                var startParse = ContentState < State.Parsing && _parsingTask == null;
                 if (startParse) {
                     Parse();
                 }
             }
-
-            if (startParse) {
-                Services.GetService<IPythonAnalyzer>().InvalidateAnalysis(this);
-            }
         }
 
-        private void LoadContent(string content) {
+        private void LoadContent(string content, int version) {
             if (ContentState < State.Loading) {
                 try {
                     content = content ?? LoadContent();
-                    _buffer.Reset(0, content);
+                    _buffer.Reset(version, content);
                     ContentState = State.Loaded;
                 } catch (IOException) { } catch (UnauthorizedAccessException) { }
             }
@@ -330,9 +313,12 @@ public void Update(IEnumerable<DocumentChange> changes) {
         public void Reset(string content) {
             lock (AnalysisLock) {
                 if (content != Content) {
-                    InitializeContent(content);
+                    ContentState = State.None;
+                    InitializeContent(content, _buffer.Version + 1);
                 }
             }
+
+            Services.GetService<IPythonAnalyzer>().InvalidateAnalysis(this);
         }
 
         private void Parse() {

src/Analysis/Ast/Impl/Modules/PythonPackage.cs

@@ -75,7 +75,5 @@ public PythonVariableModule(IPythonModule module) {
         public IMember CreateInstance(string typeName = null, LocationInfo location = null, IArgumentSet args = null) => this;
 
         public bool Equals(IPythonModule other) => other is PythonVariableModule module && Name.EqualsOrdinal(module.Name);
-
-        public Task LoadAndAnalyzeAsync(CancellationToken cancellationToken = default) => throw new NotSupportedException("Can't analyze analysis-only type");
     }
 }

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

@@ -90,13 +90,10 @@ protected override IPythonModule CreateModule(string name) {
                 return null;
             }
 
-            IPythonModule module;
-            if (!string.IsNullOrEmpty(moduleImport.ModulePath) && Uri.TryCreate(moduleImport.ModulePath, UriKind.Absolute, out var uri)) {
-                module = GetRdt().GetDocument(uri);
-                if (module != null) {
-                    GetRdt().LockDocument(uri);
-                    return module;
-                }
+            var module = GetRdt().GetDocument(moduleImport.FullName);
+            if (module != null) {
+                GetRdt().LockDocument(module.Uri);
+                return module;
             }
 
             // If there is a stub, make sure it is loaded and attached
@@ -177,7 +174,8 @@ public IPythonModule GetSpecializedModule(string name)
         => _specialized.TryGetValue(name, out var module) ? module : null;
 
         internal async Task LoadBuiltinTypesAsync(CancellationToken cancellationToken = default) {
-            await BuiltinsModule.LoadAndAnalyzeAsync(cancellationToken);
+            var analyzer = _services.GetService<IPythonAnalyzer>();
+            await analyzer.GetAnalysisAsync(BuiltinsModule, -1, cancellationToken);
 
             Check.InvalidOperation(!(BuiltinsModule.Analysis is EmptyAnalysis), "After await");
 

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

@@ -153,8 +153,6 @@ public IPythonModule GetOrCreate(string name, ModuleResolutionBase mrb) {
                 }
 
                 var module = mrb.CreateModule(name);
-                ((IDocument)module)?.Reset(null);
-
                 lock (_syncObj) {
                     _creating = false;
                     _module = module;

src/Analysis/Ast/Impl/Types/Definitions/IPythonModule.cs

@@ -47,15 +47,7 @@ public interface IPythonModule : IPythonType, IPythonFile, ILocatedMember {
         /// Global cope of the module.
         /// </summary>
         IGlobalScope GlobalScope { get; }
-
-        /// <summary>
-        /// Ensures that module content is loaded and analysis has completed.
-        /// Typically module content is loaded at the creation time, but delay
-        /// loaded (lazy) modules may choose to defer content retrieval and
-        /// analysis until later time, when module members are actually needed.
-        /// </summary>
-        Task LoadAndAnalyzeAsync(CancellationToken cancellationToken = default);
-
+
         /// <summary>
         /// If module is a stub points to the primary module.
         /// Typically used in code navigation scenarios when user