Implement declaration of submodules in module members

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

@@ -13,8 +13,12 @@
 // See the Apache Version 2.0 License for specific language governing
 // permissions and limitations under the License.
 
+using System.Collections.Generic;
+using System.Linq;
+using Microsoft.Python.Analysis.Core.DependencyResolution;
 using Microsoft.Python.Analysis.Modules;
 using Microsoft.Python.Analysis.Types;
+using Microsoft.Python.Core;
 using Microsoft.Python.Core.Text;
 using Microsoft.Python.Parsing.Ast;
 
@@ -70,5 +74,39 @@ internal static string GetComment(this IPythonModule module, int lineNum) {
 
         internal static bool IsNonUserFile(this IPythonModule module) => module.ModuleType.IsNonUserFile();
         internal static bool IsCompiled(this IPythonModule module) => module.ModuleType.IsCompiled();
+
+        public static IEnumerable<string> GetSubmoduleNames(this IPythonModule module) {
+            var searchResult = module.Interpreter.ModuleResolution.CurrentPathResolver.FindImports(module.FilePath, Enumerable.Repeat(module.Name, 1), 0, true);
+            if (searchResult is IImportChildrenSource children) {
+                foreach (var name in children.GetChildrenNames()) {
+                    if (children.TryGetChildImport(name, out var imports)) {
+                        switch (imports) {
+                            case ImplicitPackageImport packageImport:
+                                yield return packageImport.Name;
+                                break;
+                            case ModuleImport moduleImport when !moduleImport.ModulePath.PathEquals(module.FilePath):
+                                yield return moduleImport.Name;
+                                break;
+                        }
+                    }
+                }
+            }
+        }
+
+        public static IPythonModule GetSubmodule(this IPythonModule module, string subModuleName) {
+            var mres = module.Interpreter.ModuleResolution;
+            var searchResult = mres.CurrentPathResolver.FindImports(module.FilePath, Enumerable.Repeat(module.Name, 1), 0, true);
+            if (searchResult is IImportChildrenSource children) {
+                if (children.TryGetChildImport(subModuleName, out var imports)) {
+                    switch (imports) {
+                        case ImplicitPackageImport packageImport:
+                            return mres.GetImportedModule(packageImport.FullName);
+                        case ModuleImport moduleImport when !moduleImport.ModulePath.PathEquals(module.FilePath):
+                            return mres.GetImportedModule(moduleImport.FullName);
+                    }
+                }
+            }
+            return null;
+        }
     }
 }

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

@@ -22,6 +22,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Python.Analysis.Analyzer;
+using Microsoft.Python.Analysis.Core.DependencyResolution;
 using Microsoft.Python.Analysis.Dependencies;
 using Microsoft.Python.Analysis.Diagnostics;
 using Microsoft.Python.Analysis.Documents;
@@ -157,7 +158,9 @@ public virtual string Documentation {
         #endregion
 
         #region IMemberContainer
-        public virtual IMember GetMember(string name) => GlobalScope.Variables[name]?.Value;
+        public virtual IMember GetMember(string name)
+            => GlobalScope.Variables[name]?.Value ?? this.GetSubmodule(name);
+
         public virtual IEnumerable<string> GetMemberNames() {
             // drop imported modules and typing.
             return GlobalScope.Variables
@@ -166,26 +169,32 @@ public virtual IEnumerable<string> GetMemberNames() {
                     if (v.Value is IPythonInstance) {
                         return true;
                     }
+
                     var valueType = v.Value?.GetPythonType();
                     if (valueType is PythonModule) {
                         return false; // Do not re-export modules.
                     }
+
                     if (valueType is IPythonFunctionType f && f.IsLambda()) {
                         return false;
                     }
+
                     if (this is TypingModule) {
                         return true; // Let typing module behave normally.
                     }
+
                     // Do not re-export types from typing. However, do export variables
                     // assigned with types from typing. Example:
                     //    from typing import Any # do NOT export Any
                     //    x = Union[int, str] # DO export x
                     if (valueType?.DeclaringModule is TypingModule && v.Name == valueType.Name) {
                         return false;
                     }
+
                     return true;
                 })
-                .Select(v => v.Name);
+                .Select(v => v.Name)
+                .Concat(this.GetSubmoduleNames());
         }
         #endregion
 

src/Analysis/Ast/Test/ImportTests.cs

@@ -13,10 +13,13 @@
 // See the Apache Version 2.0 License for specific language governing
 // permissions and limitations under the License.
 
+using System.IO;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
 using Microsoft.Python.Analysis.Diagnostics;
+using Microsoft.Python.Analysis.Documents;
 using Microsoft.Python.Analysis.Modules;
 using Microsoft.Python.Analysis.Tests.FluentAssertions;
 using Microsoft.Python.Analysis.Types;
@@ -234,5 +237,19 @@ def exit():
             var analysis = await GetAnalysisAsync(code);
             analysis.Should().HaveVariable("x").OfType(BuiltinTypeId.Int);
         }
+
+        [TestMethod, Priority(0)]
+        public async Task ModuleMembers() {
+            var appUri = TestData.GetTestSpecificUri("app.py");
+            await TestData.CreateTestSpecificFileAsync(Path.Combine("package", "__init__.py"), "import m1");
+            await TestData.CreateTestSpecificFileAsync(Path.Combine("package", "m1", "__init__.py"), string.Empty);
+
+            await CreateServicesAsync(PythonVersions.LatestAvailable3X);
+            var rdt = Services.GetService<IRunningDocumentTable>();
+            var doc = rdt.OpenDocument(appUri, "import package");
+
+            var analysis = await doc.GetAnalysisAsync(Timeout.Infinite);
+            analysis.Should().HaveVariable("package").Which.Should().HaveMember("m1");
+        }
     }
 }

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

@@ -13,9 +13,8 @@
 // See the Apache Version 2.0 License for specific language governing
 // permissions and limitations under the License.
 
-using System.Collections.Generic;
-using System.Linq;
 using Microsoft.Python.Parsing.Ast;
+using System.Collections.Generic;
 
 namespace Microsoft.Python.Analysis.Core.DependencyResolution {
     public static class AstUtilities {

src/LanguageServer/Impl/Completion/ImportCompletion.cs

@@ -15,7 +15,6 @@
 
 using System;
 using System.Collections.Generic;
-using System.IO;
 using System.Linq;
 using Microsoft.Python.Analysis.Core.DependencyResolution;
 using Microsoft.Python.Analysis.Types;
@@ -41,7 +40,7 @@ public static CompletionResult TryGetCompletions(ImportStatement import, Complet
                 if (name != null && context.Position >= name.StartIndex) {
                     if (context.Position > name.EndIndex && name.EndIndex > name.StartIndex) {
                         var applicableSpan = context.GetApplicableSpanFromLastToken(import);
-                        return new CompletionResult(new []{ CompletionItemSource.AsKeyword }, applicableSpan);
+                        return new CompletionResult(new[] { CompletionItemSource.AsKeyword }, applicableSpan);
                     }
 
                     if (name.Names.Count == 0 || name.Names[0].EndIndex >= context.Position) {
@@ -60,7 +59,7 @@ public static CompletionResult TryGetCompletions(ImportStatement import, Complet
 
         public static CompletionResult GetCompletionsInFromImport(FromImportStatement fromImport, CompletionContext context) {
             // No more completions after '*', ever!
-            if (fromImport.Names != null && fromImport.Names.Any(n => n?.Name == "*" && context.Position > n.EndIndex)) {
+            if (fromImport.Names.Any(n => n?.Name == "*" && context.Position > n.EndIndex)) {
                 return CompletionResult.Empty;
             }
 
@@ -160,31 +159,31 @@ private static CompletionResult GetResultFromImportSearch(IImportSearchResult im
                 default:
                     return CompletionResult.Empty;
             }
-            
+
             var completions = new List<CompletionItem>();
             if (prependStar) {
                 completions.Add(CompletionItemSource.Star);
             }
 
             if (module != null) {
                 completions.AddRange(module.GetMemberNames()
-                    .Where(n => !string.IsNullOrEmpty(n))
-                    .Select(n => context.ItemSource.CreateCompletionItem(n, module.GetMember(n))));
-            }
-
-            if (importSearchResult is IImportChildrenSource children) {
-                foreach (var childName in children.GetChildrenNames()) {
-                    if (!children.TryGetChildImport(childName, out var imports)) {
-                        continue;
-                    }
-
-                    switch (imports) {
-                        case ImplicitPackageImport packageImport:
-                            completions.Add(CompletionItemSource.CreateCompletionItem(packageImport.Name, CompletionItemKind.Module));
-                            break;
-                        case ModuleImport moduleImport when !moduleImport.ModulePath.PathEquals(document.FilePath):
-                            completions.Add(CompletionItemSource.CreateCompletionItem(moduleImport.Name, CompletionItemKind.Module));
-                            break;
+                        .Where(n => !string.IsNullOrEmpty(n))
+                        .Select(n => context.ItemSource.CreateCompletionItem(n, module.GetMember(n))));
+            } else {
+                if (importSearchResult is IImportChildrenSource children) {
+                    foreach (var childName in children.GetChildrenNames()) {
+                        if (!children.TryGetChildImport(childName, out var imports)) {
+                            continue;
+                        }
+
+                        switch (imports) {
+                            case ImplicitPackageImport packageImport:
+                                completions.Add(CompletionItemSource.CreateCompletionItem(packageImport.Name, CompletionItemKind.Module));
+                                break;
+                            case ModuleImport moduleImport when !moduleImport.ModulePath.PathEquals(document.FilePath):
+                                completions.Add(CompletionItemSource.CreateCompletionItem(moduleImport.Name, CompletionItemKind.Module));
+                                break;
+                        }
                     }
                 }
             }

src/LanguageServer/Test/CompletionTests.cs

@@ -16,6 +16,7 @@
 using System;
 using System.IO;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
 using Microsoft.Python.Analysis;
@@ -1012,7 +1013,7 @@ public async Task FromDotInRootWithInitPy() {
 
             var cs = new CompletionSource(new PlainTextDocumentationSource(), ServerSettings.completion);
             var result = cs.GetCompletions(analysis, new SourceLocation(1, 7));
-            result.Should().OnlyHaveLabels("__dict__", "__file__", "__doc__", "__package__", "__debug__", "__name__", "__path__", "__spec__");
+            result.Should().OnlyHaveLabels("module1", "__dict__", "__file__", "__doc__", "__package__", "__debug__", "__name__", "__path__", "__spec__");
         }
 
         [TestMethod, Priority(0)]
@@ -1084,6 +1085,25 @@ public async Task FromDotInImplicitPackage() {
             result.Should().OnlyHaveLabels("module2", "sub_package");
         }
 
+        [TestMethod, Priority(0)]
+        public async Task SubmoduleMember() {
+            var appUri = TestData.GetTestSpecificUri("app.py");
+            await TestData.CreateTestSpecificFileAsync(Path.Combine("package", "__init__.py"), "import m1\nx = 1");
+            await TestData.CreateTestSpecificFileAsync(Path.Combine("package", "m1", "__init__.py"), string.Empty);
+            await TestData.CreateTestSpecificFileAsync(Path.Combine("package", "m2", "__init__.py"), string.Empty);
+
+            await CreateServicesAsync(PythonVersions.LatestAvailable3X);
+            var rdt = Services.GetService<IRunningDocumentTable>();
+            var doc = rdt.OpenDocument(appUri, "import package\npackage.");
+
+            await Services.GetService<IPythonAnalyzer>().WaitForCompleteAnalysisAsync();
+            var analysis = await doc.GetAnalysisAsync(Timeout.Infinite);
+            var cs = new CompletionSource(new PlainTextDocumentationSource(), ServerSettings.completion);
+
+            var result = cs.GetCompletions(analysis, new SourceLocation(2, 9));
+            result.Should().HaveLabels("m1", "m2", "x");
+        }
+
         [DataRow(false)]
         [DataRow(true)]
         [DataTestMethod, Priority(0)]