Mikhail Arkhipov (MikhailArkhipov) · Apr 30, 2019 · Mar 1, 2019 · Mar 1, 2019 · Mar 5, 2019 · Mar 5, 2019
diff --git a/src/Analysis/Ast/Impl/Analyzer/Expressions/ExpressionFinder.cs b/src/Analysis/Ast/Impl/Analyzer/Expressions/ExpressionFinder.cs
@@ -98,7 +98,7 @@ private bool Save(Node node, bool baseWalk, bool ifTrue) {
                 }
 
                 if (baseWalk && ifTrue) {
-                    Expression = node is ModuleName m && m.Names != null
+                    Expression = node is ModuleName m
                         ? m.Names.FirstOrDefault(n => n.StartIndex <= _endLocation && _endLocation <= n.EndIndex)
                         : node;
                 }

diff --git a/src/Analysis/Ast/Test/AnalysisTestBase.cs b/src/Analysis/Ast/Test/AnalysisTestBase.cs
@@ -57,6 +57,9 @@ protected virtual ServiceManager CreateServiceManager() {
 
         protected string GetAnalysisTestDataFilesPath() => TestData.GetPath(Path.Combine("TestData", "AstAnalysis"));
 
+        protected Task<IServiceManager> CreateServicesAsync(InterpreterConfiguration configuration, IServiceManager sm = null)
+            => CreateServicesAsync(TestData.GetTestSpecificRootUri().AbsolutePath, configuration, sm);
+
         protected async Task<IServiceManager> CreateServicesAsync(string root, InterpreterConfiguration configuration, IServiceManager sm = null) {
             configuration = configuration ?? PythonVersions.LatestAvailable;
             configuration.AssertInstalled();

diff --git a/src/LanguageServer/Impl/Sources/DefinitionSource.cs b/src/LanguageServer/Impl/Sources/DefinitionSource.cs
@@ -14,10 +14,10 @@
 // permissions and limitations under the License.
 
 using System;
-using System.Linq;
 using Microsoft.Python.Analysis;
 using Microsoft.Python.Analysis.Analyzer;
 using Microsoft.Python.Analysis.Analyzer.Expressions;
+using Microsoft.Python.Analysis.Core.DependencyResolution;
 using Microsoft.Python.Analysis.Documents;
 using Microsoft.Python.Analysis.Modules;
 using Microsoft.Python.Analysis.Types;
@@ -36,8 +36,8 @@ public DefinitionSource(IServiceContainer services) {
             _services = services;
         }
 
-        public Reference FindDefinition(IDocumentAnalysis analysis, SourceLocation location, out ILocatedMember member) {
-            member = null;
+        public Reference FindDefinition(IDocumentAnalysis analysis, SourceLocation location, out ILocatedMember definingMember) {
+            definingMember = null;
             if (analysis?.Ast == null) {
                 return null;
             }
@@ -46,71 +46,146 @@ public Reference FindDefinition(IDocumentAnalysis analysis, SourceLocation locat
                 FindExpressionOptions.Hover, out var exprNode, out var statement, out var exprScope);
 
             if (exprNode is ConstantExpression || !(exprNode is Expression expr)) {
-                return null; // No hover for literals.
+                return null; // No goto definition for literals.
+            }
+
+            Reference reference = null;
+            switch (statement) {
+                // Check if this is a relative import
+                case FromImportStatement fromImport:
+                    reference = HandleFromImport(analysis, location, fromImport, exprNode, out definingMember);
+                    break;
+                case ImportStatement import:
+                    reference = HandleImport(analysis, import, exprNode, out definingMember);
+                    break;
+            }
+
+            if (reference != null) {
+                return reference.uri == null ? null : reference;
             }
 
             var eval = analysis.ExpressionEvaluator;
             using (eval.OpenScope(analysis.Document, exprScope)) {
                 if (expr is MemberExpression mex) {
-                    return FromMemberExpression(mex, analysis, out member);
+                    return FromMemberExpression(mex, analysis, out definingMember);
                 }
 
                 // Try variables
                 var name = (expr as NameExpression)?.Name;
-                IMember value = null;
                 if (!string.IsNullOrEmpty(name)) {
-                    var reference = TryFromVariable(name, analysis, location, statement, out member);
-                    if (reference != null) {
-                        return reference;
-                    }
+                    reference = TryFromVariable(name, analysis, location, statement, out definingMember);
+                }
+            }
 
-                    if (statement is ImportStatement || statement is FromImportStatement) {
-                        reference = TryFromImport(statement, name, analysis, out value);
-                        if (reference != null) {
-                            member = value as ILocatedMember;
-                            return reference;
-                        }
-                    }
+            return reference;
+        }
+
+        private Reference HandleFromImport(IDocumentAnalysis analysis, SourceLocation location, FromImportStatement statement, Node expr, out ILocatedMember definingMember) {
+            definingMember = null;
+
+            // Are in the dotted name?
+            var locationIndex = location.ToIndex(analysis.Ast);
+            if (statement.Root.StartIndex <= locationIndex && locationIndex <= statement.Root.EndIndex) {
+                var mres = analysis.Document.Interpreter.ModuleResolution;
+                var imports = mres.CurrentPathResolver.FindImports(analysis.Document.FilePath, statement);
+                IPythonModule module = null;
+                switch (imports) {
+                    case ModuleImport moduleImport:
+                        module = mres.GetImportedModule(moduleImport.FullName);
+                        break;
+                    case ImplicitPackageImport packageImport:
+                        module = mres.GetImportedModule(packageImport.FullName);
+                        break;
                 }
 
-                value = value ?? eval.GetValueFromExpression(expr);
-                if (value.IsUnknown()) {
-                    return null;
+                definingMember = module;
+                return module != null
+                    ? new Reference { range = default, uri = CanNavigateToModule(module) ? module.Uri : null }
+                    : null;
+            }
+
+            // We are in what/as part
+            var nex = expr as NameExpression;
+            var name = nex?.Name;
+            if (string.IsNullOrEmpty(name)) {
+                return null;
+            }
+
+            // From X import A
+            var value = analysis.ExpressionEvaluator.GetValueFromExpression(nex);
+            if (value.IsUnknown()) {
+                // From X import A as B
+                var index = statement.Names.IndexOf(x => x?.Name == name);
+                if (index >= 0 && index < statement.AsNames.Count) {
+                    value = analysis.ExpressionEvaluator.GetValueFromExpression(statement.AsNames[index]);
                 }
-                member = value as ILocatedMember;
-                return FromMember(value);
             }
+
+            if (!value.IsUnknown()) {
+                definingMember = value as ILocatedMember;
+                return FromMember(definingMember);
+            }
+
+            return null;
         }
 
-        private Reference TryFromVariable(string name, IDocumentAnalysis analysis, SourceLocation location, Node statement, out ILocatedMember member) {
-            member = null;
+        private Reference HandleImport(IDocumentAnalysis analysis, ImportStatement statement, Node expr, out ILocatedMember definingMember) {
+            definingMember = null;
 
-            var m = analysis.ExpressionEvaluator.LookupNameInScopes(name, out var scope);
-            if (m != null && scope.Variables[name] is IVariable v) {
-                member = v;
-                var definition = v.Definition;
-                if (statement is ImportStatement || statement is FromImportStatement) {
-                    // If we are on the variable definition in this module,
-                    // then goto definition should go to the parent, if any.
-                    var indexSpan = v.Definition.Span.ToIndexSpan(analysis.Ast);
-                    var index = location.ToIndex(analysis.Ast);
-                    if (indexSpan.Start <= index && index < indexSpan.End) {
-                        if (v.Parent == null) {
-                            return null;
-                        }
-                        definition = v.Parent.Definition;
-                    }
-                }
+            var name = (expr as NameExpression)?.Name;
+            if (string.IsNullOrEmpty(name)) {
+                return null;
+            }
 
-                if (CanNavigateToModule(definition.DocumentUri)) {
-                    return new Reference { range = definition.Span, uri = definition.DocumentUri };
+            var index = statement.Names.IndexOf(x => x?.MakeString() == name);
+            if (index < 0) {
+                return null;
+            }
+
+            var module = analysis.Document.Interpreter.ModuleResolution.GetImportedModule(name);
+            if (module != null) {
+                definingMember = module;
+                return new Reference { range = default, uri = CanNavigateToModule(module) ? module.Uri : null };
+            }
+
+            // Import A as B
+            if (index >= 0 && index < statement.AsNames.Count) {
+                var value = analysis.ExpressionEvaluator.GetValueFromExpression(statement.AsNames[index]);
+                if (!value.IsUnknown()) {
+                    definingMember = value as ILocatedMember;
+                    return FromMember(definingMember);
                 }
             }
             return null;
         }
 
-        private Reference FromMemberExpression(MemberExpression mex, IDocumentAnalysis analysis, out ILocatedMember member) {
-            member = null;
+
+        private Reference TryFromVariable(string name, IDocumentAnalysis analysis, SourceLocation location, Node statement, out ILocatedMember definingMember) {
+            definingMember = null;
+
+            var m = analysis.ExpressionEvaluator.LookupNameInScopes(name, out var scope);
+            if (m == null || !(scope.Variables[name] is IVariable v)) {
+                return null;
+            }
+
+            definingMember = v;
+            if (statement is ImportStatement || statement is FromImportStatement) {
+                // If we are on the variable definition in this module,
+                // then goto definition should go to the parent, if any.
+                var indexSpan = v.Definition.Span.ToIndexSpan(analysis.Ast);
+                var index = location.ToIndex(analysis.Ast);
+                if (indexSpan.Start <= index && index < indexSpan.End) {
+                    var definition = v.Parent != null ? v.Parent.Definition : (v.Value as ILocatedMember)?.Definition;
+                    if (definition != null && CanNavigateToModule(definition.DocumentUri)) {
+                        return new Reference { range = definition.Span, uri = definition.DocumentUri };
+                    }
+                }
+            }
+            return FromMember(v);
+        }
+
+        private Reference FromMemberExpression(MemberExpression mex, IDocumentAnalysis analysis, out ILocatedMember definingMember) {
+            definingMember = null;
 
             var eval = analysis.ExpressionEvaluator;
             var target = eval.GetValueFromExpression(mex.Target);
@@ -122,7 +197,7 @@ private Reference FromMemberExpression(MemberExpression mex, IDocumentAnalysis a
                     // want the variable itself since we want to know its location.
                     var v1 = m.Analysis.GlobalScope.Variables[mex.Name];
                     if (v1 != null) {
-                        member = v1;
+                        definingMember = v1;
                         return FromMember(v1);
                     }
                     break;
@@ -133,15 +208,15 @@ private Reference FromMemberExpression(MemberExpression mex, IDocumentAnalysis a
                     using (eval.OpenScope(analysis.Document, cls.ClassDefinition)) {
                         eval.LookupNameInScopes(mex.Name, out _, out var v2, LookupOptions.Local);
                         if (v2 != null) {
-                            member = v2;
+                            definingMember = v2;
                             return FromMember(v2);
                         }
                     }
                     break;
 
                 default:
                     if (type?.GetMember(mex.Name) is ILocatedMember lm) {
-                        member = lm;
+                        definingMember = lm;
                         return FromMember(lm);
                     }
 
@@ -150,50 +225,6 @@ private Reference FromMemberExpression(MemberExpression mex, IDocumentAnalysis a
             return null;
         }
 
-        private static Reference TryFromImport(Node statement, string name, IDocumentAnalysis analysis, out IMember value) {
-            value = null;
-            string moduleName = null;
-            switch (statement) {
-                // In 'import A as B' A is not declared as a variable, so try locating B.
-                case ImportStatement imp when imp.Names.Any(x => x?.MakeString() == name):
-                case FromImportStatement fimp when fimp.Root.Names.Any(x => x?.Name == name):
-                    moduleName = name;
-                    break;
-            }
-
-            if (moduleName != null) {
-                var module = analysis.Document.Interpreter.ModuleResolution.GetImportedModule(moduleName);
-                if (module != null) {
-                    value = module;
-                    return CanNavigateToModule(module) ? new Reference { range = default, uri = module.Uri } : null;
-                }
-            }
-
-            // Perhaps it is a member such as A in 'from X import A as B'
-            switch (statement) {
-                case ImportStatement imp: {
-                        // Import A as B
-                        var index = imp.Names.IndexOf(x => x?.MakeString() == name);
-                        if (index >= 0 && index < imp.AsNames.Count) {
-                            value = analysis.ExpressionEvaluator.GetValueFromExpression(imp.AsNames[index]);
-                            return null;
-                        }
-                        break;
-                    }
-                case FromImportStatement fimp: {
-                        // From X import A as B
-                        var index = fimp.Names.IndexOf(x => x?.Name == name);
-                        if (index >= 0 && index < fimp.AsNames.Count) {
-                            value = analysis.ExpressionEvaluator.GetValueFromExpression(fimp.AsNames[index]);
-                            return null;
-                        }
-                        break;
-                    }
-            }
-            return null;
-        }
-
-
         private Reference FromMember(IMember m) {
             var definition = (m as ILocatedMember)?.Definition;
             var moduleUri = definition?.DocumentUri;
@@ -218,16 +249,9 @@ private bool CanNavigateToModule(Uri uri) {
             return CanNavigateToModule(doc);
         }
 
-        private static bool CanNavigateToModule(IPythonModule m) {
-            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.
-            canNavigate |= m.ModuleType == ModuleType.Stub || m.ModuleType == ModuleType.Compiled;
-#endif
-            return canNavigate;
-        }
+        private static bool CanNavigateToModule(IPythonModule m)
+            => m?.ModuleType == ModuleType.User ||
+               m?.ModuleType == ModuleType.Package ||
+               m?.ModuleType == ModuleType.Library;
     }
 }
diff --git a/src/LanguageServer/Test/GoToDefinitionTests.cs b/src/LanguageServer/Test/GoToDefinitionTests.cs
@@ -219,8 +219,7 @@ public async Task GotoRelativeImportInExplicitPackage() {
             var subpkgPath = TestData.GetTestSpecificUri("pkg", "subpkg", "__init__.py");
             var submodPath = TestData.GetTestSpecificUri("pkg", "submod", "__init__.py");
 
-            var root = TestData.GetTestSpecificRootUri().AbsolutePath;
-            await CreateServicesAsync(root, PythonVersions.LatestAvailable3X);
+            await CreateServicesAsync(PythonVersions.LatestAvailable3X);
             var rdt = Services.GetService<IRunningDocumentTable>();
 
             rdt.OpenDocument(pkgPath, string.Empty);
@@ -235,6 +234,34 @@ public async Task GotoRelativeImportInExplicitPackage() {
             reference.uri.Should().Be(modPath);
         }
 
+        [TestMethod, Priority(0)]
+        public async Task GotoModuleSourceFromRelativeImport() {
+            var modAPath = TestData.GetTestSpecificUri("a.py");
+            var modBPath = TestData.GetTestSpecificUri("b.py");
+
+            await CreateServicesAsync(PythonVersions.LatestAvailable3X);
+            var rdt = Services.GetService<IRunningDocumentTable>();
+
+            var modA = rdt.OpenDocument(modAPath, "from .b import X");
+            var modB = rdt.OpenDocument(modBPath, @"
+z = 1
+def X(): ...
+");
+            var analysis = await modA.GetAnalysisAsync(-1);
+            var ds = new DefinitionSource(Services);
+
+            var reference = ds.FindDefinition(analysis, new SourceLocation(1, 7), out _);
+            reference.Should().NotBeNull();
+            reference.uri.AbsolutePath.Should().Contain("b.py");
+            reference.range.Should().Be(0, 0, 0, 0);
+
+            reference = ds.FindDefinition(analysis, new SourceLocation(1, 16), out _);
+            reference.Should().NotBeNull();
+            reference.range.Should().Be(2, 4, 2, 5);
+            reference.uri.AbsolutePath.Should().Contain("b.py");
+        }
+
+
         [TestMethod, Priority(0)]
         public async Task EmptyAnalysis() {
             var analysis = await GetAnalysisAsync(string.Empty);