microsoft · MikhailArkhipov · Feb 19, 2019 · Feb 18, 2019 · Feb 19, 2019 · Feb 19, 2019
diff --git a/src/Analysis/Ast/Impl/Analyzer/Handlers/FromImportHandler.cs b/src/Analysis/Ast/Impl/Analyzer/Handlers/FromImportHandler.cs
@@ -16,6 +16,7 @@
 using System.Diagnostics;
 using System.Linq;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Python.Analysis.Core.DependencyResolution;
 using Microsoft.Python.Analysis.Modules;
 using Microsoft.Python.Analysis.Types;
@@ -42,7 +43,7 @@ public bool HandleFromImport(FromImportStatement node, CancellationToken cancell
             var imports = ModuleResolution.CurrentPathResolver.FindImports(Module.FilePath, node);
             switch (imports) {
                 case ModuleImport moduleImport when moduleImport.FullName == Module.Name:
-                    ImportMembersFromSelf(node, cancellationToken);
+                    ImportMembersFromSelf(node);
                     break;
                 case ModuleImport moduleImport:
                     ImportMembersFromModule(node, moduleImport.FullName, cancellationToken);
@@ -60,7 +61,7 @@ public bool HandleFromImport(FromImportStatement node, CancellationToken cancell
             return false;
         }
 
-        private void ImportMembersFromSelf(FromImportStatement node, CancellationToken cancellationToken = default) {
+        private void ImportMembersFromSelf(FromImportStatement node) {
             var names = node.Names;
             var asNames = node.AsNames;
 

diff --git a/src/Analysis/Ast/Impl/Diagnostics/IDiagnosticsService.cs b/src/Analysis/Ast/Impl/Diagnostics/IDiagnosticsService.cs
@@ -29,7 +29,7 @@ public interface IDiagnosticsService {
         void Replace(Uri documentUri, IEnumerable<DiagnosticsEntry> entries);
 
         /// <summary>
-        /// Removes document from the diagnostics report. Typically when document closes.
+        /// Removes document from the diagnostics report. Typically when document is disposed.
         /// </summary>
         void Remove(Uri documentUri);
 

diff --git a/src/Analysis/Ast/Impl/Modules/Resolution/MainModuleResolution.cs b/src/Analysis/Ast/Impl/Modules/Resolution/MainModuleResolution.cs
@@ -27,6 +27,7 @@
 using Microsoft.Python.Analysis.Documents;
 using Microsoft.Python.Analysis.Types;
 using Microsoft.Python.Core;
+using Microsoft.Python.Core.IO;
 using Microsoft.Python.Core.Diagnostics;
 
 namespace Microsoft.Python.Analysis.Modules.Resolution {
@@ -78,32 +79,39 @@ protected override IPythonModule CreateModule(string name) {
                 return null;
             }
 
+            var rdt = _services.GetService<IRunningDocumentTable>();
+            IPythonModule module;
+            if (!string.IsNullOrEmpty(moduleImport.ModulePath) && Uri.TryCreate(moduleImport.ModulePath, UriKind.Absolute, out var uri)) {
+                module = rdt.GetDocument(uri);
+                if (module != null) {
+                    return module;
+                }
+            }
+
             // If there is a stub, make sure it is loaded and attached
             // 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);
             }
 
-            // If stub is created and its path equals to module, return stub instead of module
-            if (stub != null && stub.FilePath.PathEquals(moduleImport.ModulePath)) {
-                return stub;
-            }
-
-            IPythonModule module;
             if (moduleImport.IsBuiltin) {
                 _log?.Log(TraceEventType.Verbose, "Create built-in compiled (scraped) module: ", name, Configuration.InterpreterPath);
                 module = new CompiledBuiltinPythonModule(name, stub, _services);
             } else if (moduleImport.IsCompiled) {
                 _log?.Log(TraceEventType.Verbose, "Create compiled (scraped): ", moduleImport.FullName, moduleImport.ModulePath, moduleImport.RootPath);
                 module = new CompiledPythonModule(moduleImport.FullName, ModuleType.Compiled, moduleImport.ModulePath, stub, _services);
             } else {
-                _log?.Log(TraceEventType.Verbose, "Create: ", moduleImport.FullName, moduleImport.ModulePath);
-                var rdt = _services.GetService<IRunningDocumentTable>();
-                // TODO: handle user code and library module separately.
+                _log?.Log(TraceEventType.Verbose, "Import: ", moduleImport.FullName, moduleImport.ModulePath);
+                // Module inside workspace == user code.
+
+                var moduleType = moduleImport.ModulePath.IsUnderRoot(_root, _fs.StringComparison)
+                    ? ModuleType.User
+                    : ModuleType.Library;
+
                 var mco = new ModuleCreationOptions {
                     ModuleName = moduleImport.FullName,
-                    ModuleType = ModuleType.Library,
+                    ModuleType = moduleType,
                     FilePath = moduleImport.ModulePath,
                     Stub = stub
                 };
@@ -203,5 +211,8 @@ private bool TryCreateModuleStub(string name, string modulePath, out IPythonModu
             module = !string.IsNullOrEmpty(stubPath) ? new StubPythonModule(name, stubPath, false, _services) : null;
             return module != null;
         }
+
+        protected override void ReportModuleNotFound(string name)
+            => _log?.Log(TraceEventType.Information, $"Import not found: {name}");
     }
 }
diff --git a/src/Analysis/Ast/Impl/Modules/Resolution/ModuleResolutionBase.cs b/src/Analysis/Ast/Impl/Modules/Resolution/ModuleResolutionBase.cs
@@ -159,5 +159,7 @@ public IPythonModule GetOrCreate(string name, ModuleResolutionBase mrb) {
                 }
             }
         }
+
+        protected abstract void ReportModuleNotFound(string name);
     }
 }
diff --git a/src/Analysis/Ast/Impl/Modules/Resolution/TypeshedResolution.cs b/src/Analysis/Ast/Impl/Modules/Resolution/TypeshedResolution.cs
@@ -139,5 +139,7 @@ private bool IsPackage(string directory)
                 !string.IsNullOrEmpty(ModulePath.GetPackageInitPy(directory)) :
                 _fs.DirectoryExists(directory);
 
+        protected override void ReportModuleNotFound(string name)
+            => _log?.Log(TraceEventType.Verbose, $"Typeshed stub not found for '{name}'");
     }
 }
diff --git a/src/Core/Impl/Extensions/IOExtensions.cs b/src/Core/Impl/Extensions/IOExtensions.cs
@@ -124,5 +124,13 @@ public static void WriteTextWithRetry(this IFileSystem fs, string filePath, stri
                 fs.DeleteFile(filePath);
             } catch (IOException) { } catch (UnauthorizedAccessException) { }
         }
+
+        public static bool IsUnderRoot(this string path, string root, StringComparison comparison) {
+            if (!string.IsNullOrEmpty(root)) {
+                var normalized = PathUtils.NormalizePath(path);
+                return normalized.StartsWith(root, comparison);
+            }
+            return false;
+        }
     }
 }
diff --git a/src/LanguageServer/Impl/Diagnostics/DiagnosticsService.cs b/src/LanguageServer/Impl/Diagnostics/DiagnosticsService.cs
@@ -17,6 +17,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using Microsoft.Python.Analysis.Diagnostics;
+using Microsoft.Python.Analysis.Documents;
 using Microsoft.Python.Core;
 using Microsoft.Python.Core.Disposables;
 using Microsoft.Python.Core.Idle;
@@ -26,17 +27,40 @@
 
 namespace Microsoft.Python.LanguageServer.Diagnostics {
     internal sealed class DiagnosticsService : IDiagnosticsService, IDisposable {
-        private readonly Dictionary<Uri, List<DiagnosticsEntry>> _diagnostics = new Dictionary<Uri, List<DiagnosticsEntry>>();
+        private sealed class DocumentDiagnostics {
+            private DiagnosticsEntry[] _entries;
+
+            public DiagnosticsEntry[] Entries {
+                get => _entries ?? Array.Empty<DiagnosticsEntry>();
+                set {
+                    _entries = value;
+                    Changed = true;
+                }
+            }
+            public bool Changed { get; set; }
+
+            public void Clear() {
+                Changed = _entries.Length > 0;
+                _entries = Array.Empty<DiagnosticsEntry>();
+            }
+        }
+
+        private readonly Dictionary<Uri, DocumentDiagnostics> _diagnostics = new Dictionary<Uri, DocumentDiagnostics>();
         private readonly DisposableBag _disposables = DisposableBag.Create<DiagnosticsService>();
+        private readonly IServiceContainer _services;
         private readonly IClientApplication _clientApp;
         private readonly object _lock = new object();
         private DiagnosticsSeverityMap _severityMap = new DiagnosticsSeverityMap();
+        private IRunningDocumentTable _rdt;
         private DateTime _lastChangeTime;
-        private bool _changed;
+
+        private IRunningDocumentTable Rdt => _rdt ?? (_rdt = _services.GetService<IRunningDocumentTable>());
 
         public DiagnosticsService(IServiceContainer services) {
-            var idleTimeService = services.GetService<IIdleTimeService>();
+            _services = services;
+            _clientApp = services.GetService<IClientApplication>();
 
+            var idleTimeService = services.GetService<IIdleTimeService>();
             if (idleTimeService != null) {
                 idleTimeService.Idle += OnIdle;
                 idleTimeService.Closing += OnClosing;
@@ -45,40 +69,38 @@ public DiagnosticsService(IServiceContainer services) {
                     .Add(() => idleTimeService.Idle -= OnIdle)
                     .Add(() => idleTimeService.Idle -= OnClosing);
             }
-            _clientApp = services.GetService<IClientApplication>();
         }
 
         #region IDiagnosticsService
         public IReadOnlyDictionary<Uri, IReadOnlyList<DiagnosticsEntry>> Diagnostics {
             get {
                 lock (_lock) {
                     return _diagnostics.ToDictionary(kvp => kvp.Key,
-                        kvp => kvp.Value
-                            .Where(e => DiagnosticsSeverityMap.GetEffectiveSeverity(e.ErrorCode, e.Severity) != Severity.Suppressed)
-                            .Select(e => new DiagnosticsEntry(
-                                e.Message,
-                                e.SourceSpan,
-                                e.ErrorCode,
-                                DiagnosticsSeverityMap.GetEffectiveSeverity(e.ErrorCode, e.Severity))
-                            ).ToList() as IReadOnlyList<DiagnosticsEntry>);
+                        kvp => FilterBySeverityMap(kvp.Value).ToList() as IReadOnlyList<DiagnosticsEntry>);
                 }
             }
         }
 
         public void Replace(Uri documentUri, IEnumerable<DiagnosticsEntry> entries) {
             lock (_lock) {
-                _diagnostics[documentUri] = entries.ToList();
+                if (!_diagnostics.TryGetValue(documentUri, out var documentDiagnostics)) {
+                    documentDiagnostics = new DocumentDiagnostics();
+                    _diagnostics[documentUri] = documentDiagnostics;
+                }
+                documentDiagnostics.Entries = entries.ToArray();
+                documentDiagnostics.Changed = true;
                 _lastChangeTime = DateTime.Now;
-                _changed = true;
             }
         }
 
         public void Remove(Uri documentUri) {
             lock (_lock) {
                 // Before removing the document, make sure we clear its diagnostics.
-                _diagnostics[documentUri] = new List<DiagnosticsEntry>();
-                PublishDiagnostics();
-                _diagnostics.Remove(documentUri);
+                if (_diagnostics.TryGetValue(documentUri, out var d)) {
+                    d.Clear();
+                    PublishDiagnostics();
+                    _diagnostics.Remove(documentUri);
+                }
             }
         }
 
@@ -89,6 +111,10 @@ public DiagnosticsSeverityMap DiagnosticsSeverityMap {
             set {
                 lock (_lock) {
                     _severityMap = value;
+                    foreach (var d in _diagnostics) {
+                        _diagnostics[d.Key].Changed = true;
+                        _lastChangeTime = DateTime.Now;
+                    }
                     PublishDiagnostics();
                 }
             }
@@ -103,22 +129,29 @@ public void Dispose() {
         private void OnClosing(object sender, EventArgs e) => Dispose();
 
         private void OnIdle(object sender, EventArgs e) {
-            if (_changed && (DateTime.Now - _lastChangeTime).TotalMilliseconds > PublishingDelay) {
+            if ((DateTime.Now - _lastChangeTime).TotalMilliseconds > PublishingDelay) {
+                ConnectToRdt();
                 PublishDiagnostics();
             }
         }
 
         private void PublishDiagnostics() {
-            KeyValuePair<Uri, IReadOnlyList<DiagnosticsEntry>>[] diagnostics;
+            var diagnostics = new List<KeyValuePair<Uri, DocumentDiagnostics>>();
             lock (_lock) {
-                diagnostics = Diagnostics.ToArray();
-                _changed = false;
+                foreach (var d in _diagnostics) {
+                    if (d.Value.Changed) {
+                        diagnostics.Add(d);
+                        _diagnostics[d.Key].Changed = false;
+                    }
+                }
             }
 
             foreach (var kvp in diagnostics) {
                 var parameters = new PublishDiagnosticsParams {
                     uri = kvp.Key,
-                    diagnostics = kvp.Value.Select(ToDiagnostic).ToArray()
+                    diagnostics = Rdt.GetDocument(kvp.Key)?.IsOpen == true
+                            ? FilterBySeverityMap(kvp.Value).Select(ToDiagnostic).ToArray()
+                            : Array.Empty<Diagnostic>()
                 };
                 _clientApp.NotifyWithParameterObjectAsync("textDocument/publishDiagnostics", parameters).DoNotWait();
             }
@@ -127,7 +160,6 @@ private void PublishDiagnostics() {
         private void ClearAllDiagnostics() {
             lock (_lock) {
                 _diagnostics.Clear();
-                _changed = false;
             }
         }
 
@@ -156,5 +188,48 @@ private static Diagnostic ToDiagnostic(DiagnosticsEntry e) {
                 message = e.Message,
             };
         }
+
+        private IEnumerable<DiagnosticsEntry> FilterBySeverityMap(DocumentDiagnostics d)
+           => d.Entries
+                .Where(e => DiagnosticsSeverityMap.GetEffectiveSeverity(e.ErrorCode, e.Severity) != Severity.Suppressed)
+                .Select(e => new DiagnosticsEntry(
+                    e.Message,
+                    e.SourceSpan,
+                    e.ErrorCode,
+                    DiagnosticsSeverityMap.GetEffectiveSeverity(e.ErrorCode, e.Severity))
+                );
+
+        private void ConnectToRdt() {
+            if (_rdt == null) {
+                _rdt = _services.GetService<IRunningDocumentTable>();
+                if (_rdt != null) {
+                    _rdt.Opened += OnOpenDocument;
+                    _rdt.Closed += OnCloseDocument;
+
+                    _disposables
+                        .Add(() => _rdt.Opened -= OnOpenDocument)
+                        .Add(() => _rdt.Closed -= OnCloseDocument);
+                }
+            }
+        }
+
+        private void OnOpenDocument(object sender, DocumentEventArgs e) {
+            lock (_lock) {
+                if(_diagnostics.TryGetValue(e.Document.Uri, out var d)) {
+                    d.Changed = d.Entries.Length > 0;
+                }
+            }
+        }
+
+        private void OnCloseDocument(object sender, DocumentEventArgs e) {
+            lock (_lock) {
+                // Before removing the document, make sure we clear its diagnostics.
+                if (_diagnostics.TryGetValue(e.Document.Uri, out var d)) {
+                    d.Clear();
+                    PublishDiagnostics();
+                    _diagnostics.Remove(e.Document.Uri);
+                }
+            }
+        }
     }
 }
diff --git a/src/LanguageServer/Impl/Implementation/Server.Documents.cs b/src/LanguageServer/Impl/Implementation/Server.Documents.cs
@@ -16,13 +16,11 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Python.Analysis;
 using Microsoft.Python.Analysis.Documents;
 using Microsoft.Python.Core;
-using Microsoft.Python.Core.Text;
 using Microsoft.Python.LanguageServer.Protocol;
 
 namespace Microsoft.Python.LanguageServer.Implementation {
-Original file line number
+Diff line change
@@ Expand Up @@
                     }
                 }
             }
+            protected abstract void ReportModuleNotFound(string name);
         }
     }