Fix goto definition in stubbed modules (microsoft#659)

* Merge part I

* Merge part II

* Fix nav to builtins

Author: Mikhail Arkhipov

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

@@ -78,8 +78,11 @@ public IMember LookupNameInScopes(string name, out IScope scope, LookupOptions o
             scope = scopes.FirstOrDefault(s => s.Variables.Contains(name));
             var value = scope?.Variables[name].Value;
             if (value == null) {
-                if (Module != Interpreter.ModuleResolution.BuiltinsModule && options.HasFlag(LookupOptions.Builtins)) {
-                    value = Interpreter.ModuleResolution.BuiltinsModule.GetMember(name);
+                var builtins = Interpreter.ModuleResolution.BuiltinsModule;
+                value = Interpreter.ModuleResolution.BuiltinsModule.GetMember(name);
+                if (Module != builtins && options.HasFlag(LookupOptions.Builtins)) {
+                    value = builtins.GetMember(name);
+                    scope = builtins.GlobalScope;
                 }
             }
 
@@ -100,7 +103,7 @@ public IPythonType GetTypeFromAnnotation(Expression expr, LookupOptions options
                 // Try generics
                 var target = GetValueFromExpression(indexExpr.Target);
                 var result = GetValueFromGeneric(target, indexExpr);
-                if(result != null) {
+                if (result != null) {
                     return result.GetPythonType();
                 }
             }

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

@@ -108,15 +108,15 @@ private void MergeStub() {
 
                         // Get documentation from the current type, if any, since stubs
                         // typically do not contain documentation while scraped code does.
-                        member?.GetPythonType()?.TransferDocumentation(stubMember.GetPythonType());
+                        member?.GetPythonType()?.TransferDocumentationAndLocation(stubMember.GetPythonType());
                         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.IsUnknown() && !(stubType is IPythonModule)) {
-                        sourceType.TransferDocumentation(stubType);
+                        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, LocationInfo.Empty, overwrite: true);

src/Analysis/Ast/Impl/Extensions/PythonTypeExtensions.cs

@@ -21,18 +21,23 @@ public static class PythonTypeExtensions {
         public static bool IsUnknown(this IPythonType value) =>
             value == null || (value.TypeId == BuiltinTypeId.Unknown && value.MemberType == PythonMemberType.Unknown && value.Name.Equals("Unknown"));
 
-        public static bool IsGenericParameter(this IPythonType value) 
+        public static bool IsGenericParameter(this IPythonType value)
             => value is IGenericTypeParameter;
 
         public static bool IsGeneric(this IPythonType value)
             => value is IGenericTypeParameter || value is IGenericType || (value is IPythonClassType c && c.IsGeneric());
 
-        public static void TransferDocumentation(this IPythonType src, IPythonType dst) {
+        public static void TransferDocumentationAndLocation(this IPythonType src, IPythonType dst) {
             if (src != null && dst is PythonType pt) {
                 pt.TrySetTypeId(dst.TypeId);
                 var documentation = src.Documentation;
-                if (string.IsNullOrEmpty(pt.Documentation) && !string.IsNullOrEmpty(documentation)) {
-                    pt.SetDocumentationProvider(_ => documentation);
+                if (!string.IsNullOrEmpty(documentation)) {
+                    if (!string.IsNullOrEmpty(documentation)) {
+                        pt.SetDocumentation(documentation);
+                    }
+                }
+                if (src is ILocatedMember lm) {
+                    pt.SetLocation(lm.Location);
                 }
             }
         }

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

@@ -103,6 +103,9 @@ internal PythonModule(ModuleCreationOptions creationOptions, IServiceContainer s
             Uri = uri;
             FilePath = creationOptions.FilePath ?? uri?.LocalPath;
             Stub = creationOptions.Stub;
+            if (Stub is PythonModule stub && ModuleType != ModuleType.Stub) {
+                stub.PrimaryModule = this;
+            }
 
             if (ModuleType == ModuleType.Specialized || ModuleType == ModuleType.Unresolved) {
                 ContentState = State.Analyzed;
@@ -199,6 +202,13 @@ public async Task LoadAndAnalyzeAsync(CancellationToken cancellationToken = defa
             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
+        /// wants to see library code and not a stub.
+        /// </summary>
+        public IPythonModule PrimaryModule { get; internal set; }
+
         protected virtual string LoadContent() {
             if (ContentState < State.Loading) {
                 ContentState = State.Loading;

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

@@ -17,7 +17,7 @@ namespace Microsoft.Python.Analysis.Types {
     /// <summary>
     /// Represents member of a class.
     /// </summary>
-    public interface IPythonClassMember : IPythonType {
+    public interface IPythonClassMember : IPythonType, ILocatedMember {
         IPythonType DeclaringType { get; }
     }
 }

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

@@ -20,7 +20,7 @@ namespace Microsoft.Python.Analysis.Types {
     /// <summary>
     /// Represents Python class type definition.
     /// </summary>
-    public interface IPythonClassType : IPythonType {
+    public interface IPythonClassType : IPythonType, ILocatedMember {
         /// <summary>
         /// Class definition node in the AST.
         /// </summary>

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

@@ -55,5 +55,12 @@ public interface IPythonModule : IPythonType, IPythonFile, ILocatedMember {
         /// 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
+        /// wants to see library code and not a stub.
+        /// </summary>
+        IPythonModule PrimaryModule { get; }
     }
 }

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

@@ -116,7 +116,7 @@ internal override void SetDocumentationProvider(Func<string, string> provider) {
         #region IPythonFunction
         public FunctionDefinition FunctionDefinition { get; }
         public IPythonType DeclaringType { get; }
-        public override string Documentation => _doc ?? _overloads.FirstOrDefault()?.Documentation;
+        public override string Documentation => _doc ?? base.Documentation ?? _overloads.FirstOrDefault()?.Documentation;
         public virtual bool IsClassMethod { get; private set; }
         public virtual bool IsStatic { get; private set; }
         public override bool IsAbstract => _isAbstract;

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

@@ -23,9 +23,11 @@ namespace Microsoft.Python.Analysis.Types {
     [DebuggerDisplay("{Name}")]
     internal class PythonType : IPythonType, ILocatedMember, IHasQualifiedName, IEquatable<IPythonType> {
         private readonly object _lock = new object();
-        private readonly Func<string, LocationInfo> _locationProvider;
         private readonly string _name;
         private Func<string, string> _documentationProvider;
+        private Func<string, LocationInfo> _locationProvider;
+        private string _documentation;
+        private LocationInfo _location;
         private Dictionary<string, IMember> _members;
         private BuiltinTypeId _typeId;
         private bool _readonly;
@@ -41,26 +43,32 @@ public PythonType(
             string documentation,
             LocationInfo location,
             BuiltinTypeId typeId = BuiltinTypeId.Unknown
-        ) : this(name, declaringModule, _ => documentation, _ => location ?? LocationInfo.Empty, typeId) { }
+        ) : this(name, declaringModule, typeId) {
+            _documentation = documentation;
+            _location = location;
+        }
 
         public PythonType(
-            string name,
-            IPythonModule declaringModule,
-            Func<string, string> documentationProvider,
-            Func<string, LocationInfo> locationProvider,
-            BuiltinTypeId typeId = BuiltinTypeId.Unknown
-        ) {
-            _name = name ?? throw new ArgumentNullException(nameof(name));
-            DeclaringModule = declaringModule;
+                string name,
+                IPythonModule declaringModule,
+                Func<string, string> documentationProvider,
+                Func<string, LocationInfo> locationProvider,
+                BuiltinTypeId typeId = BuiltinTypeId.Unknown
+            ) : this(name, declaringModule, typeId) {
             _documentationProvider = documentationProvider;
             _locationProvider = locationProvider;
-            _typeId = typeId;
+        }
+
+        private PythonType(string name, IPythonModule declaringModule, BuiltinTypeId typeId = BuiltinTypeId.Unknown) {
+            _name = name ?? throw new ArgumentNullException(nameof(name));
+            DeclaringModule = declaringModule;
+            _typeId = typeId; _typeId = typeId;
         }
 
         #region IPythonType
 
         public virtual string Name => TypeId == BuiltinTypeId.Ellipsis ? "..." : _name;
-        public virtual string Documentation => _documentationProvider?.Invoke(Name);
+        public virtual string Documentation => _documentationProvider != null ? _documentationProvider.Invoke(Name) : _documentation;
         public IPythonModule DeclaringModule { get; }
         public virtual PythonMemberType MemberType => _typeId.GetMemberId();
         public virtual BuiltinTypeId TypeId => _typeId;
@@ -96,7 +104,8 @@ public virtual IMember Call(IPythonInstance instance, string memberName, IArgume
         #endregion
 
         #region ILocatedMember
-        public virtual LocationInfo Location => _locationProvider?.Invoke(Name) ?? LocationInfo.Empty;
+        public virtual LocationInfo Location => _locationProvider != null
+            ? _locationProvider?.Invoke(Name) : _location ?? LocationInfo.Empty;
         #endregion
 
         #region IHasQualifiedName
@@ -119,6 +128,8 @@ internal bool TrySetTypeId(BuiltinTypeId typeId) {
         }
 
         internal virtual void SetDocumentationProvider(Func<string, string> provider) => _documentationProvider = provider;
+        internal virtual void SetDocumentation(string documentation) => _documentation = documentation;
+        internal virtual void SetLocation(LocationInfo location) => _location = location;
 
         internal void AddMembers(IEnumerable<IVariable> variables, bool overwrite) {
             lock (_lock) {

src/LanguageServer/Impl/Sources/DefinitionSource.cs

@@ -14,7 +14,6 @@
 // permissions and limitations under the License.
 
 using Microsoft.Python.Analysis;
-using Microsoft.Python.Analysis.Analyzer;
 using Microsoft.Python.Analysis.Analyzer.Expressions;
 using Microsoft.Python.Analysis.Documents;
 using Microsoft.Python.Analysis.Modules;
@@ -40,80 +39,106 @@ public Reference FindDefinition(IDocumentAnalysis analysis, SourceLocation locat
 
             var eval = analysis.ExpressionEvaluator;
             using (eval.OpenScope(analysis.Document, exprScope)) {
+
+                // First try variables, except in imports
+                if (expr is NameExpression nex && !string.IsNullOrEmpty(nex.Name) &&
+                    !(statement is ImportStatement) && !(statement is FromImportStatement)) {
+                    var m = eval.LookupNameInScopes(nex.Name, out var scope);
+                    if (m != null && scope.Variables[nex.Name] is IVariable v) {
+                        if (CanNavigateToModule(v.Value.GetPythonType()?.DeclaringModule, analysis)) {
+                            return new Reference { range = v.Location.Span, uri = v.Location.DocumentUri };
+                        }
+                    }
+                }
+
                 var value = eval.GetValueFromExpression(expr);
-                return FromMember(value, expr, eval);
+                return FromMember(value, expr, statement, analysis);
             }
         }
 
-        private Reference FromMember(IMember value, Expression expr, IExpressionEvaluator eval) {
+        private Reference FromMember(IMember value, Expression expr, Node statement, IDocumentAnalysis analysis) {
             Node node = null;
             IPythonModule module = null;
+            LocationInfo location = null;
+            var eval = analysis.ExpressionEvaluator;
 
             switch (value) {
                 case IPythonClassType cls:
                     node = cls.ClassDefinition;
                     module = cls.DeclaringModule;
+                    location = cls.Location;
                     break;
                 case IPythonFunctionType fn:
                     node = fn.FunctionDefinition;
                     module = fn.DeclaringModule;
+                    location = fn.Location;
                     break;
                 case IPythonPropertyType prop:
                     node = prop.FunctionDefinition;
                     module = prop.DeclaringModule;
+                    location = prop.Location;
                     break;
-                case IPythonModule mod: {
-                        var member = eval.LookupNameInScopes(mod.Name, out var scope);
-                        if (member != null && scope != null) {
-                            var v = scope.Variables[mod.Name];
-                            if (v != null) {
-                                return new Reference { range = v.Location.Span, uri = v.Location.DocumentUri };
-                            }
-                        }
-                        break;
-                    }
-                case IPythonInstance instance when instance.Type is IPythonFunctionType ft: {
-                        node = ft.FunctionDefinition;
-                        module = ft.DeclaringModule;
-                        break;
-                    }
-                case IPythonInstance _ when expr is NameExpression nex: {
-                        var member = eval.LookupNameInScopes(nex.Name, out var scope);
-                        if (member != null && scope != null) {
-                            var v = scope.Variables[nex.Name];
-                            if (v != null) {
-                                return new Reference { range = v.Location.Span, uri = v.Location.DocumentUri };
-                            }
+                case IPythonModule mod:
+                    return HandleModule(mod, analysis, statement);
+                case IPythonInstance instance when instance.Type is IPythonFunctionType ft:
+                    node = ft.FunctionDefinition;
+                    module = ft.DeclaringModule;
+                    location = ft.Location;
+                    break;
+                case IPythonInstance instance when instance.Type is IPythonFunctionType ft:
+                    node = ft.FunctionDefinition;
+                    module = ft.DeclaringModule;
+                    break;
+                case IPythonInstance _ when expr is NameExpression nex:
+                    var m1 = eval.LookupNameInScopes(nex.Name, out var scope);
+                    if (m1 != null && scope != null) {
+                        var v = scope.Variables[nex.Name];
+                        if (v != null) {
+                            return new Reference { range = v.Location.Span, uri = v.Location.DocumentUri };
                         }
-                        break;
                     }
+                    break;
                 case IPythonInstance _ when expr is MemberExpression mex: {
                         var target = eval.GetValueFromExpression(mex.Target);
                         var type = target?.GetPythonType();
-                        var member = type?.GetMember(mex.Name);
-                        if (member is IPythonInstance v) {
+                        var m2 = type?.GetMember(mex.Name);
+                        if (m2 is IPythonInstance v) {
                             return new Reference { range = v.Location.Span, uri = v.Location.DocumentUri };
                         }
-                        return FromMember(member, null, eval);
+                        return FromMember(m2, null, statement, analysis);
                     }
             }
 
-            if (node != null && CanNavigateToModule(module) && module is IDocument doc) {
+            module = module?.ModuleType == ModuleType.Stub ? module.PrimaryModule : module;
+            if (node != null && module is IDocument doc && CanNavigateToModule(module, analysis)) {
                 return new Reference {
-                    range = node.GetSpan(doc.GetAnyAst()), uri = doc.Uri
+                    range = location?.Span ?? node.GetSpan(doc.GetAnyAst()), uri = doc.Uri
                 };
             }
 
             return null;
         }
 
-        private static bool CanNavigateToModule(IPythonModule m) {
+        private static Reference HandleModule(IPythonModule module, IDocumentAnalysis analysis, Node statement) {
+            // If we are in import statement, open the module source if available.
+            if (statement is ImportStatement || statement is FromImportStatement) {
+                if (module.Uri != null && CanNavigateToModule(module, analysis)) {
+                    return new Reference { range = default, uri = module.Uri };
+                }
+            }
+            return null;
+        }
+
+        private static bool CanNavigateToModule(IPythonModule m, IDocumentAnalysis analysis) {
+            if (m == null) {
+                return false;
+            }
+            var canNavigate = m.ModuleType == ModuleType.User || m.ModuleType == ModuleType.Package || m.ModuleType == ModuleType.Library;
 #if DEBUG
             // Allow navigation anywhere in debug.
-            return m.ModuleType != ModuleType.Specialized && m.ModuleType != ModuleType.Unresolved; 
-#else
-            return m.ModuleType == ModuleType.User || m.ModuleType == ModuleType.Package || m.ModuleType == ModuleType.Library;
+            canNavigate |= m.ModuleType == ModuleType.Stub || m.ModuleType == ModuleType.Compiled;
 #endif
+            return canNavigate;
         }
     }
 }