Fix tests

Author: MikhailArkhipov

src/Analysis/Ast/Impl/Analyzer/Handlers/FromImportHandler.cs

@@ -40,8 +40,16 @@ public bool HandleFromImport(FromImportStatement node) {
                 }
             }
 
-            FindModule(node.Root, null, node.ForceAbsolute, out var firstModule, out var lastModule, out _, out var imports);
-            AssignVariables(node, imports, lastModule);
+            if(node.Root is RelativeModuleName) {
+                var imports = ModuleResolution.CurrentPathResolver.FindImports(Module.FilePath, node);
+                if (HandleImportSearchResult(imports, null, null, node.Root, out var variableModule)) {
+                    AssignVariables(node, imports, variableModule);
+                }
+            } else {
+                FindModuleByAbsoluteName(node.Root, null, node.ForceAbsolute, out _, out var module, out _, out var imports);
+                AssignVariables(node, imports, module);
+            }
+
             return false;
         }
 

src/Analysis/Ast/Impl/Analyzer/Handlers/ImportHandler.cs

@@ -49,7 +49,19 @@ public bool HandleImport(ImportStatement node) {
             return false;
         }
 
-        private void FindModule(ModuleName moduleImportExpression, NameExpression asNameExpression, bool forceAbsolute, 
+        private void HandleImport(ModuleName moduleImportExpression, NameExpression asNameExpression, bool forceAbsolute) {
+            FindModuleByAbsoluteName(moduleImportExpression, asNameExpression, forceAbsolute, out var firstModule, out var lastModule, out var importNames, out _);
+            // "import fob.oar.baz as baz" is handled as baz = import_module('fob.oar.baz')
+            // "import fob.oar.baz" is handled as fob = import_module('fob')
+            if (!string.IsNullOrEmpty(asNameExpression?.Name) && lastModule != default) {
+                Eval.DeclareVariable(asNameExpression.Name, lastModule, VariableSource.Import, asNameExpression);
+            } else if (firstModule != default && !string.IsNullOrEmpty(importNames[0])) {
+                var firstName = moduleImportExpression.Names[0];
+                Eval.DeclareVariable(importNames[0], firstModule, VariableSource.Import, firstName);
+            }
+        }
+
+        private void FindModuleByAbsoluteName(ModuleName moduleImportExpression, NameExpression asNameExpression, bool forceAbsolute,
             out PythonVariableModule firstModule, out PythonVariableModule lastModule, out ImmutableArray<string> importNames, out IImportSearchResult imports) {
             // "import fob.oar.baz" means
             // import_module('fob')
@@ -73,18 +85,6 @@ private void FindModule(ModuleName moduleImportExpression, NameExpression asName
             }
         }
 
-        private void HandleImport(ModuleName moduleImportExpression, NameExpression asNameExpression, bool forceAbsolute) {
-            FindModule(moduleImportExpression, asNameExpression, forceAbsolute, out var firstModule, out var lastModule, out var importNames, out _);
-            // "import fob.oar.baz as baz" is handled as baz = import_module('fob.oar.baz')
-            // "import fob.oar.baz" is handled as fob = import_module('fob')
-            if (!string.IsNullOrEmpty(asNameExpression?.Name) && lastModule != default) {
-                Eval.DeclareVariable(asNameExpression.Name, lastModule, VariableSource.Import, asNameExpression);
-            } else if (firstModule != default && !string.IsNullOrEmpty(importNames[0])) {
-                var firstName = moduleImportExpression.Names[0];
-                Eval.DeclareVariable(importNames[0], firstModule, VariableSource.Import, firstName);
-            }
-        }
-
         private bool HandleImportSearchResult(in IImportSearchResult imports, in PythonVariableModule parent, in NameExpression asNameExpression, in Node location, out PythonVariableModule variableModule) {
             switch (imports) {
                 case ModuleImport moduleImport when Module.ModuleType == ModuleType.Stub && moduleImport.FullName == Module.Name:

src/Analysis/Core/Impl/DependencyResolution/PathResolverExtensions.cs

@@ -0,0 +1,33 @@
+// Copyright(c) Microsoft Corporation
+// All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the License); you may not use
+// this file except in compliance with the License. You may obtain a copy of the
+// License at http://www.apache.org/licenses/LICENSE-2.0
+//
+// THIS CODE IS PROVIDED ON AN  *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS
+// OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY
+// IMPLIED WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
+// MERCHANTABILITY OR NON-INFRINGEMENT.
+//
+// See the Apache Version 2.0 License for specific language governing
+// permissions and limitations under the License.
+
+using Microsoft.Python.Parsing.Ast;
+
+namespace Microsoft.Python.Analysis.Core.DependencyResolution {
+    public static class PathResolverExtensions {
+        /// <summary>
+        /// Given module file path and module name expression locates the module.
+        /// Module name expression can be absolute or relative. <see cref="ModuleName"/>
+        /// and <see cref="RelativeModuleName"/>.
+        /// </summary>
+        /// <returns></returns>
+        public static IImportSearchResult FindImports(this PathResolverSnapshot pathResolver, string modulePath, ModuleName moduleImportExpression, bool forceAbsolute) {
+            var rootNames = moduleImportExpression.Names.Select(n => n.Name);
+            return moduleImportExpression is RelativeModuleName relativeName
+                ? pathResolver.GetImportsFromRelativePath(modulePath, relativeName.DotCount, rootNames)
+                : pathResolver.GetImportsFromAbsoluteName(modulePath, rootNames, forceAbsolute);
+        }
+    }
+}