From 67fed10d1c892de07a9a920b5396d53d7473b977 Mon Sep 17 00:00:00 2001
From: MikhailArkhipov <mikhaila@microsoft.com>
Date: Wed, 24 Jul 2019 16:45:54 -0700
Subject: [PATCH] Unify package search in imports

---
 .../Impl/Analyzer/Handlers/FromImportHandler.cs   |  8 +++-----
 .../Ast/Impl/Analyzer/Handlers/ImportHandler.cs   | 15 ++++++++++-----
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/Analysis/Ast/Impl/Analyzer/Handlers/FromImportHandler.cs b/src/Analysis/Ast/Impl/Analyzer/Handlers/FromImportHandler.cs
index 182f83d5d..f4ca23b2f 100644
--- a/src/Analysis/Ast/Impl/Analyzer/Handlers/FromImportHandler.cs
+++ b/src/Analysis/Ast/Impl/Analyzer/Handlers/FromImportHandler.cs
@@ -40,10 +40,8 @@ public bool HandleFromImport(FromImportStatement node) {
                 }
             }
 
-            var imports = ModuleResolution.CurrentPathResolver.FindImports(Module.FilePath, node);
-            if (HandleImportSearchResult(imports, null, null, node.Root, out var variableModule)) {
-                AssignVariables(node, imports, variableModule);
-            }
+            FindModule(node.Root, null, node.ForceAbsolute, out var firstModule, out var lastModule, out _, out var imports);
+            AssignVariables(node, imports, lastModule);
             return false;
         }
 
@@ -82,7 +80,7 @@ private void HandleModuleImportStar(PythonVariableModule variableModule, bool is
             }
 
             // If __all__ is present, take it, otherwise declare all members from the module that do not begin with an underscore.
-            var memberNames = isImplicitPackage 
+            var memberNames = isImplicitPackage
                 ? variableModule.GetMemberNames()
                 : variableModule.Analysis.StarImportMemberNames ?? variableModule.GetMemberNames().Where(s => !s.StartsWithOrdinal("_"));
 
diff --git a/src/Analysis/Ast/Impl/Analyzer/Handlers/ImportHandler.cs b/src/Analysis/Ast/Impl/Analyzer/Handlers/ImportHandler.cs
index af71f6556..4973fd14d 100644
--- a/src/Analysis/Ast/Impl/Analyzer/Handlers/ImportHandler.cs
+++ b/src/Analysis/Ast/Impl/Analyzer/Handlers/ImportHandler.cs
@@ -49,17 +49,19 @@ public bool HandleImport(ImportStatement node) {
             return false;
         }
 
-        private void HandleImport(ModuleName moduleImportExpression, NameExpression asNameExpression, bool forceAbsolute) {
+        private void FindModule(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')
             // import_module('fob.oar')
             // import_module('fob.oar.baz')
-            var importNames = ImmutableArray<string>.Empty;
-            var lastModule = default(PythonVariableModule);
-            var firstModule = default(PythonVariableModule);
+            importNames = ImmutableArray<string>.Empty;
+            lastModule = default;
+            firstModule = default;
+            imports = null;
             foreach (var nameExpression in moduleImportExpression.Names) {
                 importNames = importNames.Add(nameExpression.Name);
-                var imports = ModuleResolution.CurrentPathResolver.GetImportsFromAbsoluteName(Module.FilePath, importNames, forceAbsolute);
+                imports = ModuleResolution.CurrentPathResolver.GetImportsFromAbsoluteName(Module.FilePath, importNames, forceAbsolute);
                 if (!HandleImportSearchResult(imports, lastModule, asNameExpression, moduleImportExpression, out lastModule)) {
                     lastModule = default;
                     break;
@@ -69,7 +71,10 @@ private void HandleImport(ModuleName moduleImportExpression, NameExpression asNa
                     firstModule = lastModule;
                 }
             }
+        }
 
+        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) {