Allow navigation to stubs (#1960)

* Remove stale reference

* Don't suppress LHS diagnostics on augmented assign

* Revert "Don't suppress LHS diagnostics on augmented assign"

This reverts commit 6109ac7.

* Escape [ and ]

* PR feedback

* Allow navigation to stubs

* Better handle stubs for compiled modules

* Fix typeshed path

* Partial undo

* Partial undo

* Undo accidental change

Author: Mikhail Arkhipov

src/Analysis/Ast/Impl/Analyzer/Evaluation/ExpressionEval.Scopes.cs

@@ -179,6 +179,8 @@ public IDisposable OpenScope(IPythonModule module, ScopeStatement node, out Scop
             return new ScopeTracker(this);
         }
 
+        internal void ReplaceVariable(IVariable v) => CurrentScope.ReplaceVariable(v);
+
         private class ScopeTracker : IDisposable {
             private readonly ExpressionEval _eval;
 

src/Analysis/Ast/Impl/Analyzer/Evaluation/ExpressionEval.cs

@@ -60,7 +60,9 @@ public ExpressionEval(IServiceContainer services, IPythonModule module, PythonAs
         public LocationInfo GetLocationInfo(Node node) => node?.GetLocation(this) ?? LocationInfo.Empty;
 
         public Location GetLocationOfName(Node node) {
-            if (node == null || (Module.ModuleType != ModuleType.User && Module.ModuleType != ModuleType.Library)) {
+            if (node == null || 
+                Module.ModuleType == ModuleType.Specialized || Module.ModuleType == ModuleType.Compiled || 
+                Module.ModuleType == ModuleType.CompiledBuiltin || Module.ModuleType == ModuleType.Builtins) {
                 return DefaultLocation;
             }
 

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

@@ -110,8 +110,16 @@ private void TryReplaceMember(IVariable v, IPythonType sourceType, IPythonType s
                     }
                     break;
 
-                case PythonClassType sourceClass:
-                    MergeClass(v, sourceClass, stubType, cancellationToken);
+                case IPythonClassType sourceClass:
+                    MergeMembers(v, sourceClass, stubType, cancellationToken);
+                    break;
+
+                case PythonFunctionType sourceFunction:
+                    MergeMembers(v, sourceFunction, stubType, cancellationToken);
+                    break;
+
+                case PythonPropertyType sourceProperty:
+                    MergeMembers(v, sourceProperty, stubType, cancellationToken);
                     break;
 
                 case IPythonModule _:
@@ -137,28 +145,38 @@ private void TryReplaceMember(IVariable v, IPythonType sourceType, IPythonType s
             }
         }
 
-        private void MergeClass(IVariable v, IPythonClassType sourceClass, IPythonType stubType, CancellationToken cancellationToken) {
+        private void MergeMembers(IVariable v, IPythonType sourceType, IPythonType stubType, CancellationToken cancellationToken) {
             // Transfer documentation first so we get class documentation
             // that comes from the class definition win over one that may
             // come from __init__ during the member merge below.
-            TransferDocumentationAndLocation(sourceClass, stubType);
+            TransferDocumentationAndLocation(sourceType.GetPythonType(), stubType);
 
             // Replace the class entirely since stub members may use generic types
             // and the class definition is important. We transfer missing members
             // from the original class to the stub.
-            _eval.DeclareVariable(v.Name, v.Value, v.Source);
+            //
+            // In case module is compiled, it is already a stub and has no locations
+            // for code navigation.. In this case we replace the entire variable by one
+            // from the stub rather than just the value since stub variable has location
+            // and its own root definition/reference chain.
+            if (sourceType.DeclaringModule.ModuleType == ModuleType.Compiled || 
+                sourceType.DeclaringModule.ModuleType == ModuleType.CompiledBuiltin) {
+                _eval.ReplaceVariable(v);
+            } else {
+                _eval.DeclareVariable(v.Name, v.Value, v.Source);
+            }
 
             // First pass: go through source class members and pick those 
             // that are not present in the stub class.
-            foreach (var name in sourceClass.GetMemberNames().ToArray()) {
+            foreach (var name in sourceType.GetMemberNames().ToArray()) {
                 cancellationToken.ThrowIfCancellationRequested();
 
-                var sourceMember = sourceClass.GetMember(name);
+                var sourceMember = sourceType.GetMember(name);
                 if (sourceMember.IsUnknown()) {
                     continue; // Do not add unknowns to the stub.
                 }
                 var sourceMemberType = sourceMember?.GetPythonType();
-                if (sourceMemberType is IPythonClassMember cm && cm.DeclaringType != sourceClass) {
+                if (sourceMemberType is IPythonClassMember cm && !cm.DeclaringModule.Equals(sourceType.DeclaringModule)) {
                     continue; // Only take members from this class and not from bases.
                 }
                 if (!IsFromThisModuleOrSubmodules(sourceMemberType)) {
@@ -196,7 +214,7 @@ private void MergeClass(IVariable v, IPythonClassType sourceClass, IPythonType s
                     continue; // Only take members from this class and not from bases.
                 }
 
-                var sourceMember = sourceClass.GetMember(name);
+                var sourceMember = sourceType.GetMember(name);
                 if (sourceMember.IsUnknown()) {
                     continue;
                 }
@@ -305,7 +323,13 @@ private bool IsFromThisModuleOrSubmodules(IPythonType type) {
             var thisModule = _eval.Module;
             var typeModule = type.DeclaringModule;
             var typeMainModuleName = typeModule.Name.Split('.').FirstOrDefault();
-            return typeModule.Equals(thisModule) || typeMainModuleName == thisModule.Name;
+            if (typeModule.Equals(thisModule) || typeMainModuleName == thisModule.Name) {
+                return true;
+            }
+            // Check if module is explicitly imported by the current one. For example, 'os'
+            // imports 'nt' and os.pyi specifies functions from 'nt' such as mkdir and so on.
+            var imported = thisModule.GlobalScope.Variables[typeModule.Name];
+            return imported?.Value != null && imported.Source == VariableSource.Import;
         }
     }
 }

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

@@ -32,8 +32,11 @@ internal sealed class TypeshedResolution : ModuleResolutionBase, IModuleResoluti
 
         public TypeshedResolution(string root, IServiceContainer services) : base(root, services) {
             // TODO: merge with user-provided stub paths
-            var stubs = Path.Combine(Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location), "Stubs");
-            _typeStubPaths = GetTypeShedPaths(Root)
+            var asmLocation = Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location);
+            var stubs = Path.Combine(asmLocation, "Stubs");
+            var typeshedRoot = Root ?? Path.Combine(asmLocation, "Typeshed");
+            
+            _typeStubPaths = GetTypeShedPaths(typeshedRoot)
                 .Concat(GetTypeShedPaths(stubs))
                 .Where(services.GetService<IFileSystem>().DirectoryExists)
                 .ToImmutableArray();

src/Analysis/Ast/Impl/Values/Scope.cs

@@ -93,6 +93,11 @@ public void DeclareImported(string name, IMember value, Location location = defa
 
         internal void AddChildScope(Scope s) => (_childScopes ?? (_childScopes = new List<Scope>())).Add(s);
 
+        internal void ReplaceVariable(IVariable v) {
+            VariableCollection.RemoveVariable(v.Name);
+            VariableCollection.DeclareVariable(v.Name, v.Value, v.Source, v.Location);
+        }
+
         private VariableCollection VariableCollection => _variables ?? (_variables = new VariableCollection());
 
         private void DeclareBuiltinVariables() {

src/LanguageServer/Impl/Implementation/Server.Editor.cs

@@ -89,17 +89,23 @@ public async Task<Reference[]> GotoDefinition(TextDocumentPositionParams @params
             _log?.Log(TraceEventType.Verbose, $"Goto Definition in {uri} at {@params.position}");
 
             var analysis = await Document.GetAnalysisAsync(uri, Services, CompletionAnalysisTimeout, cancellationToken);
-            var reference = new DefinitionSource(Services).FindDefinition(analysis, @params.position, out _);
-            return reference != null ? new[] { reference } : Array.Empty<Reference>();
+            var ds = new DefinitionSource(Services);
+            var reference = ds.FindDefinition(analysis, @params.position, out _);
+            return reference != null && ds.CanNavigateToModule(reference.uri) 
+                ? new[] { reference } 
+                : Array.Empty<Reference>();
         }
 
         public async Task<Location> GotoDeclaration(TextDocumentPositionParams @params, CancellationToken cancellationToken) {
             var uri = @params.textDocument.uri;
             _log?.Log(TraceEventType.Verbose, $"Goto Declaration in {uri} at {@params.position}");
 
             var analysis = await Document.GetAnalysisAsync(uri, Services, CompletionAnalysisTimeout, cancellationToken);
-            var reference = new DeclarationSource(Services).FindDefinition(analysis, @params.position, out _);
-            return reference != null ? new Location { uri = reference.uri, range = reference.range } : null;
+            var ds = new DeclarationSource(Services);
+            var reference = ds.FindDefinition(analysis, @params.position, out _);
+            return reference != null && ds.CanNavigateToModule(reference.uri) 
+                ? new Location { uri = reference.uri, range = reference.range } 
+                : null;
         }
 
         public Task<Reference[]> FindReferences(ReferencesParams @params, CancellationToken cancellationToken) {
@@ -137,8 +143,8 @@ public async Task<CodeAction[]> CodeAction(CodeActionParams @params, Cancellatio
 
             return codeActions.ToArray();
 
-            static bool AskedFor(CodeActionParams @params, string codeActionKind) {
-                return @params.context.only == null || @params.context.only.Any(s => s.StartsWith(codeActionKind));
+            bool AskedFor(CodeActionParams p, string codeActionKind) {
+                return p.context.only == null || p.context.only.Any(s => s.StartsWith(codeActionKind));
             }
         }
     }

src/LanguageServer/Impl/Sources/DefinitionSource.cs

@@ -14,6 +14,7 @@
 // permissions and limitations under the License.
 
 using System;
+using System.IO;
 using System.Linq;
 using Microsoft.Python.Analysis;
 using Microsoft.Python.Analysis.Analyzer;
@@ -360,22 +361,32 @@ private Reference FromMember(IMember m) {
             return null;
         }
 
-        private bool CanNavigateToModule(Uri uri) {
+        public bool CanNavigateToModule(Uri uri) {
             if (uri == null) {
                 return false;
             }
+
+            if (!CanNavigateToPath(uri.LocalPath)) {
+                return false;
+            }
             var rdt = _services.GetService<IRunningDocumentTable>();
             var doc = rdt.GetDocument(uri);
             //  Allow navigation to modules not in RDT - most probably
             // it is a module that was restored from database.
             return doc == null || CanNavigateToModule(doc);
         }
 
-        private static bool CanNavigateToModule(IPythonModule m)
-            => m?.ModuleType == ModuleType.User ||
-               m?.ModuleType == ModuleType.Stub ||
-               m?.ModuleType == ModuleType.Package ||
-               m?.ModuleType == ModuleType.Library ||
-               m?.ModuleType == ModuleType.Specialized;
+        private static bool CanNavigateToModule(IPythonModule m) {
+            if(m == null || !CanNavigateToPath(m.FilePath)) {
+                return false;
+            }
+            return m.ModuleType == ModuleType.User ||
+                m.ModuleType == ModuleType.Stub ||
+                m.ModuleType == ModuleType.Package ||
+                m.ModuleType == ModuleType.Library ||
+                m.ModuleType == ModuleType.Specialized;
+        }
+
+        private static bool CanNavigateToPath(string path) => Path.GetExtension(path) != ".exe";
     }
 }

src/LanguageServer/Test/GoToDefinitionTests.cs

@@ -720,5 +720,19 @@ def func(a, b):
             var reference = ds.FindDefinition(analysis, new SourceLocation(3, 12), out _);
             reference.range.Should().Be(1, 9, 1, 10);
         }
+
+        [TestMethod, Priority(0)]
+        public async Task CompiledCode() {
+            const string code = @"
+import os
+os.mkdir()
+";
+            var analysis = await GetAnalysisAsync(code);
+            var ds = new DefinitionSource(Services);
+            var reference = ds.FindDefinition(analysis, new SourceLocation(3, 6), out _);
+            reference.range.start.line.Should().NotBe(1);
+            reference.range.end.line.Should().NotBe(1);
+            reference.uri.AbsolutePath.Should().Contain(".pyi");
+        }
     }
 }

src/LanguageServer/Test/HoverTests.cs

@@ -284,6 +284,18 @@ class A:
             AssertHover(hs, analysis, new SourceLocation(4, 7), @"x: int", new SourceSpan(4, 7, 4, 8));
         }
 
+        [TestMethod, Priority(0)]
+        public async Task CompiledCode() {
+            const string code = @"
+import os
+os.mkdir()
+";
+            var analysis = await GetAnalysisAsync(code);
+            var hs = new HoverSource(new PlainTextDocumentationSource());
+            var hover = hs.GetHover(analysis, new SourceLocation(3, 6));
+            hover.contents.value.Should().Contain("Create a directory");
+        }
+
         private static void AssertNoHover(HoverSource hs, IDocumentAnalysis analysis, SourceLocation position) {
             var hover = hs.GetHover(analysis, position);
             hover.Should().BeNull();