fix completion of top-level packages

python/src/com/jetbrains/python/psi/impl/PyImportReferenceImpl.java

@@ -261,6 +261,9 @@ public class PyImportReferenceImpl extends PyReferenceImpl {
           }
           myObjects.addAll(Arrays.asList(completionVariants));
         }
+        else {
+          myObjects.addAll(PyModuleType.getSubmoduleVariants(target_dir, myElement, myNamesAlready));
+        }
       }
     }
   }

python/src/com/jetbrains/python/psi/types/PyModuleType.java

@@ -61,29 +61,28 @@ public class PyModuleType implements PyType { // Modules don't descend from obje
 
 
   /**
-   * @return a list of submodules of this module, either files or dirs, for easier naming; may contain filenames
+   * @param directory the module directory
+   *
+   * @return a list of submodules of the specified module directory, either files or dirs, for easier naming; may contain filenames
    *         not suitable for import.
    */
   @NotNull
-  public List<PsiFileSystemItem> getSubmodulesList() {
+  private static List<PsiFileSystemItem> getSubmodulesList(final PsiDirectory directory) {
     List<PsiFileSystemItem> result = new ArrayList<PsiFileSystemItem>();
 
-    if (PyNames.INIT_DOT_PY.equals(myModule.getName())) { // our module is a dir, not a single file
-      PsiDirectory mydir = myModule.getContainingDirectory();
-      if (mydir != null) { // just in case
-        // file modules
-        for (PsiFile f : mydir.getFiles()) {
-          final String filename = f.getName();
-          // if we have a binary module, we'll most likely also have a stub for it in site-packages
-          if ((f instanceof PyFile && !filename.equals(PyNames.INIT_DOT_PY)) || isBinaryModule(filename)) {
-            result.add(f);
-          }
-        }
-        // dir modules
-        for (PsiDirectory dir : mydir.getSubdirectories()) {
-          if (dir.findFile(PyNames.INIT_DOT_PY) instanceof PyFile) result.add(dir);
+    if (directory != null) { // just in case
+      // file modules
+      for (PsiFile f : directory.getFiles()) {
+        final String filename = f.getName();
+        // if we have a binary module, we'll most likely also have a stub for it in site-packages
+        if ((f instanceof PyFile && !filename.equals(PyNames.INIT_DOT_PY)) || isBinaryModule(filename)) {
+          result.add(f);
         }
       }
+      // dir modules
+      for (PsiDirectory dir : directory.getSubdirectories()) {
+        if (dir.findFile(PyNames.INIT_DOT_PY) instanceof PyFile) result.add(dir);
+      }
     }
     return result;
   }
@@ -128,22 +127,32 @@ public class PyModuleType implements PyType { // Modules don't descend from obje
       }
     }
     if (point == ResolveImportUtil.PointInImport.AS_MODULE || point == ResolveImportUtil.PointInImport.AS_NAME) { // when imported from somehow, add submodules
-      for (PsiFileSystemItem pfsi : getSubmodulesList()) {
-        if (pfsi == location.getContainingFile().getOriginalFile()) continue;
-        String s = pfsi.getName();
-        int pos = s.lastIndexOf('.'); // it may not contain a dot, except in extension; cut it off.
-        if (pos > 0) s = s.substring(0, pos);
-        if (!PyNames.isIdentifier(s)) continue;
-        if (names_already != null) {
-          if (names_already.contains(s)) continue;
-          else names_already.add(s);
-        }
-        result.add(LookupElementBuilder.create(pfsi, s).setPresentableText(s).setIcon(pfsi.getIcon(0)));
+      if (PyNames.INIT_DOT_PY.equals(myModule.getName())) { // our module is a dir, not a single file
+        result.addAll(getSubmoduleVariants(myModule.getContainingDirectory(), location, names_already));
       }
     }
     return result.toArray();
   }
 
+  public static List<LookupElement> getSubmoduleVariants(final PsiDirectory directory,
+                                                         PsiElement location,
+                                                         Set<String> names_already) {
+    List<LookupElement> result = new ArrayList<LookupElement>();
+    for (PsiFileSystemItem pfsi : getSubmodulesList(directory)) {
+      if (pfsi == location.getContainingFile().getOriginalFile()) continue;
+      String s = pfsi.getName();
+      int pos = s.lastIndexOf('.'); // it may not contain a dot, except in extension; cut it off.
+      if (pos > 0) s = s.substring(0, pos);
+      if (!PyNames.isIdentifier(s)) continue;
+      if (names_already != null) {
+        if (names_already.contains(s)) continue;
+        else names_already.add(s);
+      }
+      result.add(LookupElementBuilder.create(pfsi, s).setPresentableText(s).setIcon(pfsi.getIcon(0)));
+    }
+    return result;
+  }
+
   public String getName() {
     PsiFile mod = getModule();
     if (mod != null) {

python/testData/completion/import.after.py

@@ -0,0 +1 @@
+import datetime

python/testData/completion/import.py

@@ -0,0 +1 @@
+import da<caret>

python/testSrc/com/jetbrains/python/PythonCompletionTest.java

@@ -289,4 +289,8 @@ public class PythonCompletionTest extends PyLightFixtureTestCase {
     myFixture.completeBasic();
     myFixture.checkResultByFile("completion/relativeImport/pkg/name.after.py");
   }
+
+  public void testImport() {
+    doTest();
+  }
 }