PY-12497 .NET stubs are generated incorrectly when name is imported instead of module

python/helpers/generator3.py

@@ -2,9 +2,12 @@
 import atexit
 import zipfile
 
+# TODO: Move all CLR-specific functions to clr_tools
+
 from pycharm_generator_utils.module_redeclarator import *
 from pycharm_generator_utils.util_methods import *
 from pycharm_generator_utils.constants import *
+from pycharm_generator_utils.clr_tools import *
 
 
 debug_mode = False
@@ -430,6 +433,9 @@ if __name__ == "__main__":
             if '-p' in opts:
                 atexit.register(print_profile)
 
+            # We take module name from import statement
+            name = get_namespace_by_name(name)
+
         if not process_one(name, mod_file_name, False, subdir):
             sys.exit(1)
 

python/helpers/pycharm_generator_utils/clr_tools.py

@@ -0,0 +1,60 @@
+# coding=utf-8
+"""
+.NET (CLR) specific functions
+"""
+__author__ = 'Ilya.Kazakevich'
+
+
+def get_namespace_by_name(object_name):
+    """
+    Gets namespace for full object name. Sometimes last element of name is module while it may be class.
+    For System.Console returns System, for System.Web returns System.Web.
+    Be sure all required assemblies are loaded (i.e. clr.AddRef.. is called)
+    :param object_name: name to parse
+    :return: namespace
+    """
+    (imported_object, object_name) = _import_first(object_name)
+    parts = object_name.partition(".")
+    first_part = parts[0]
+    remain_part = parts[2]
+
+    while remain_part and type(_get_attr_by_name(imported_object, remain_part)) is type: # While we are in class
+        remain_part = remain_part.rpartition(".")[0]
+
+    return first_part + "." + remain_part if remain_part else first_part
+
+def _import_first(object_name):
+    """
+    Some times we can not import module directly. For example, Some.Class.InnerClass could not be imported: you need to import "Some.Class"
+    or even "Some" instead. This function tries to find part of name that could be loaded
+
+     :param object_name: name in dotted notation like "Some.Function.Here"
+     :return: (imported_object, object_name): tuple with object and its name
+    """
+    while object_name:
+        try:
+            return (__import__(object_name), object_name)
+        except ImportError:
+            object_name = object_name.rpartition(".")[0] # Remove rightest part
+    raise Exception("No module name found in name " + object_name)
+
+
+
+def _get_attr_by_name(obj, name):
+    """
+    Accepts chain of attributes in dot notation like "some.property.name" and gets them on object
+    :param obj: object to introspec
+    :param name: attribute name
+    :return attribute
+
+    >>> str(_get_attr_by_name("A", "__class__.__class__"))
+    "<type 'type'>"
+
+    >>> str(_get_attr_by_name("A", "__class__.__len__.__class__"))
+    "<type 'method_descriptor'>"
+    """
+    result = obj
+    parts = name.split('.')
+    for part in parts:
+        result = getattr(result, part)
+    return result

python/psi-api/src/com/jetbrains/python/psi/PyImportStatementBase.java

@@ -17,6 +17,8 @@ package com.jetbrains.python.psi;
 
 import org.jetbrains.annotations.NotNull;
 
+import java.util.List;
+
 /**
  * @author yole
  */
@@ -26,4 +28,11 @@ public interface PyImportStatementBase extends PyStatement {
    */
   @NotNull
   PyImportElement[] getImportElements();
+
+  /**
+   * @return qualified names of imported elements regardless way they were imported.
+   * "from bar import foo" or "import bar.foo" or "from bar import foo as spam" are all "bar.foo"
+   */
+  @NotNull
+  List<String> getFullyQualifiedObjectNames();
 }

python/src/com/jetbrains/python/PyBundle.properties

@@ -829,6 +829,7 @@ sdk.gen.reading.versions.file=Reading versions file...
 sdk.gen.notify.converting.old.skels=Converting old skeletons
 sdk.gen.notify.converting.text=Skeletons of binary modules seem to be from an older version.<br/>These will be fully re-generated, which will take some time, but will happen <i>only once</i>.<br/>Next time you open the project, only skeletons of new or updated binary modules will be re-generated.
 sdk.gen.updating.skels=Updating skeletons
+sdk.gen.stubs.for.binary.modules=Generate stubs for binary module {0}
 
 # remote interpreters
 remote.interpreter.configure.title=Configure Remote Python Interpreter

python/src/com/jetbrains/python/inspections/PyUnresolvedReferencesInspection.java

@@ -419,13 +419,14 @@ public class PyUnresolvedReferencesInspection extends PyInspection {
         return;
       }
       String description = null;
-      final String text = reference.getElement().getText();
+      PsiElement element = reference.getElement();
+      final String text = element.getText();
       TextRange rangeInElement = reference.getRangeInElement();
       String refText = text;  // text of the part we're working with
       if (rangeInElement.getStartOffset() > 0 && rangeInElement.getEndOffset() > 0) {
         refText = rangeInElement.substring(text);
       }
-      final PsiElement element = reference.getElement();
+
       final List<LocalQuickFix> actions = new ArrayList<LocalQuickFix>(2);
       final String refName = (element instanceof PyQualifiedExpression) ? ((PyQualifiedExpression)element).getReferencedName() : refText;
       // Empty text, nothing to highlight
@@ -565,8 +566,11 @@ public class PyUnresolvedReferencesInspection extends PyInspection {
         hl_type = ProblemHighlightType.LIKE_UNKNOWN_SYMBOL;
       }
 
-      if (GenerateBinaryStubsFix.isApplicable(reference)) {
-        actions.add(new GenerateBinaryStubsFix(reference));
+      if (element != null) {
+        PyImportStatementBase importStatementBase = PsiTreeUtil.getParentOfType(element, PyImportStatementBase.class);
+        if ((importStatementBase != null) && GenerateBinaryStubsFix.isApplicable(importStatementBase)) {
+          actions.addAll(GenerateBinaryStubsFix.generateFixes(importStatementBase));
+        }
       }
       if (canonicalQName != null) {
         actions.add(new AddIgnoredIdentifierQuickFix(canonicalQName, false));

python/src/com/jetbrains/python/inspections/quickfix/GenerateBinaryStubsFix.java

@@ -32,13 +32,11 @@ import com.intellij.openapi.vfs.LocalFileSystem;
 import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiReference;
-import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.Consumer;
+import com.jetbrains.python.PyBundle;
 import com.jetbrains.python.PythonHelpersLocator;
 import com.jetbrains.python.psi.*;
-import com.jetbrains.python.psi.impl.references.PyImportReference;
 import com.jetbrains.python.sdk.InvalidSdkException;
 import com.jetbrains.python.sdk.PySdkUtil;
 import com.jetbrains.python.sdk.PythonSdkType;
@@ -51,6 +49,7 @@ import org.jetbrains.annotations.Nullable;
 
 import java.io.File;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 
@@ -60,36 +59,38 @@ import java.util.List;
 public class GenerateBinaryStubsFix implements LocalQuickFix {
   private static final Logger LOG = Logger.getInstance("#" + GenerateBinaryStubsFix.class.getName());
 
-  private String myQualifiedName;
-  private Sdk mySdk;
+  private final String myQualifiedName;
+  private final Sdk mySdk;
 
-  public GenerateBinaryStubsFix(PsiReference reference) {
-    myQualifiedName = getReferenceText(reference);
-    mySdk = getPythonSdk(reference);
+  /**
+   * Generates pack of fixes available for some unresolved import statement.
+   * Be sure to call {@link #isApplicable(com.jetbrains.python.psi.PyImportStatementBase)} first to make sure this statement is supported
+   * @param importStatementBase statement to fix
+   * @return pack of fixes
+   */
+  @NotNull
+  public static Collection<GenerateBinaryStubsFix> generateFixes(@NotNull final PyImportStatementBase importStatementBase) {
+    final List<String> names = importStatementBase.getFullyQualifiedObjectNames();
+    final List<GenerateBinaryStubsFix> result = new ArrayList<GenerateBinaryStubsFix>(names.size());
+    for (final String qualifiedName : names) {
+        result.add(new GenerateBinaryStubsFix(importStatementBase, qualifiedName));
+    }
+    return result;
   }
 
-  private static String getReferenceText(PsiReference reference) {
-    PsiElement element = reference.getElement();
-    while (element.getParent() instanceof PyReferenceExpression) {
-      element = element.getParent();
-    }
-    final String elementText = element.getText();
-
-    final PyFromImportStatement importStatementBase = PsiTreeUtil.getParentOfType(element, PyFromImportStatement.class);
-    if (importStatementBase != null) {
-      final QualifiedName qName = importStatementBase.getImportSourceQName();
-      if (qName != null) {
-        return qName.append(elementText).toString();
-      }
-    }
-
-    return elementText;
+  /**
+   * @param importStatementBase statement to fix
+   * @param qualifiedName name should be fixed (one of {@link com.jetbrains.python.psi.PyImportStatementBase#getFullyQualifiedObjectNames()})
+   */
+  private GenerateBinaryStubsFix(@NotNull final PyImportStatementBase importStatementBase, @NotNull final String qualifiedName) {
+    myQualifiedName = qualifiedName;
+    mySdk = getPythonSdk(importStatementBase);
   }
 
   @Override
   @NotNull
   public String getName() {
-    return "Generate stubs for binary module " + myQualifiedName;
+    return PyBundle.message("sdk.gen.stubs.for.binary.modules", myQualifiedName);
   }
 
   @Override
@@ -181,7 +182,8 @@ public class GenerateBinaryStubsFix implements LocalQuickFix {
       @Override
       public void visitPyCallExpression(PyCallExpression node) {
         super.visitPyCallExpression(node);
-        if (node.isCalleeText("AddReference", "AddReferenceByPartialName")) {
+        // TODO: What if user loads it not by literal? We need to ask user for list of DLLs
+        if (node.isCalleeText("AddReference", "AddReferenceByPartialName", "AddReferenceByName")) {
           final PyExpression[] args = node.getArguments();
           if (args.length == 1 && args [0] instanceof PyStringLiteralExpression) {
             result.add(((PyStringLiteralExpression) args [0]).getStringValue());
@@ -192,11 +194,17 @@ public class GenerateBinaryStubsFix implements LocalQuickFix {
     return result;
   }
 
-  public static boolean isApplicable(@NotNull final PsiReference ref) {
-    if (!(ref instanceof PyImportReference)) {
+  /**
+   * Checks if this fix can help you to generate binary stubs
+   *
+   * @param importStatementBase statement to fix
+   * @return true if this fix could work
+   */
+  public static boolean isApplicable(@NotNull final PyImportStatementBase importStatementBase) {
+    if (importStatementBase.getFullyQualifiedObjectNames().isEmpty()) {
       return false;
     }
-    final Sdk sdk = getPythonSdk(ref);
+    final Sdk sdk = getPythonSdk(importStatementBase);
     if (sdk == null) {
       return false;
     }
@@ -204,22 +212,22 @@ public class GenerateBinaryStubsFix implements LocalQuickFix {
     if (flavor instanceof IronPythonSdkFlavor) {
       return true;
     }
-    return isGtk(ref);
+    return isGtk(importStatementBase);
   }
 
-  private static boolean isGtk(@NotNull final PsiReference ref) {
-    final PyFromImportStatement importStatementBase = PsiTreeUtil.getParentOfType(ref.getElement(), PyFromImportStatement.class);
-    if (importStatementBase != null) {
-      final QualifiedName qName = importStatementBase.getImportSourceQName();
-      if (qName != null && qName.matches("gi", "repository"))
+  private static boolean isGtk(@NotNull final PyImportStatementBase importStatementBase) {
+    if (importStatementBase instanceof PyFromImportStatement) {
+      final QualifiedName qName = ((PyFromImportStatement)importStatementBase).getImportSourceQName();
+      if (qName != null && qName.matches("gi", "repository")) {
         return true;
+      }
     }
     return false;
   }
 
   @Nullable
-  private static Sdk getPythonSdk(@NotNull final PsiReference ref) {
-    final Module module = ModuleUtilCore.findModuleForPsiElement(ref.getElement());
-    return module == null ? null : PythonSdkType.findPythonSdk(module);
+  private static Sdk getPythonSdk(@NotNull final PsiElement element) {
+    final Module module = ModuleUtilCore.findModuleForPsiElement(element);
+    return (module == null) ? null : PythonSdkType.findPythonSdk(module);
   }
 }

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

@@ -100,9 +100,9 @@ public class PyFromImportStatementImpl extends PyBaseElementImpl<PyFromImportSta
     List<PyImportElement> result = new ArrayList<PyImportElement>();
     final ASTNode importKeyword = getNode().findChildByType(PyTokenTypes.IMPORT_KEYWORD);
     if (importKeyword != null) {
-      for(ASTNode node = importKeyword.getTreeNext(); node != null; node = node.getTreeNext()) {
+      for (ASTNode node = importKeyword.getTreeNext(); node != null; node = node.getTreeNext()) {
         if (node.getElementType() == PyElementTypes.IMPORT_ELEMENT) {
-          result.add((PyImportElement) node.getPsi());
+          result.add((PyImportElement)node.getPsi());
         }
       }
     }
@@ -121,7 +121,7 @@ public class PyFromImportStatementImpl extends PyBaseElementImpl<PyFromImportSta
 
     int result = 0;
     ASTNode seeker = getNode().getFirstChildNode();
-    while(seeker != null && (seeker.getElementType() == PyTokenTypes.FROM_KEYWORD || seeker.getElementType() == TokenType.WHITE_SPACE)) {
+    while (seeker != null && (seeker.getElementType() == PyTokenTypes.FROM_KEYWORD || seeker.getElementType() == TokenType.WHITE_SPACE)) {
       seeker = seeker.getTreeNext();
     }
     while (seeker != null && seeker.getElementType() == PyTokenTypes.DOT) {
@@ -146,7 +146,9 @@ public class PyFromImportStatementImpl extends PyBaseElementImpl<PyFromImportSta
     return findChildByType(PyTokenTypes.RPAR);
   }
 
-  public boolean processDeclarations(@NotNull final PsiScopeProcessor processor, @NotNull final ResolveState state, final PsiElement lastParent,
+  public boolean processDeclarations(@NotNull final PsiScopeProcessor processor,
+                                     @NotNull final ResolveState state,
+                                     final PsiElement lastParent,
                                      @NotNull final PsiElement place) {
     // import is per-file
     if (place.getContainingFile() != getContainingFile()) {
@@ -165,7 +167,7 @@ public class PyFromImportStatementImpl extends PyBaseElementImpl<PyFromImportSta
     }
     else {
       PyImportElement[] importElements = getImportElements();
-      for(PyImportElement element: importElements) {
+      for (PyImportElement element : importElements) {
         if (!processor.execute(element, state)) {
           return false;
         }
@@ -192,7 +194,8 @@ public class PyFromImportStatementImpl extends PyBaseElementImpl<PyFromImportSta
     ASTNode prevNode = result;
     do {
       prevNode = prevNode.getTreePrev();
-    } while(prevNode != null && prevNode.getElementType() == TokenType.WHITE_SPACE);
+    }
+    while (prevNode != null && prevNode.getElementType() == TokenType.WHITE_SPACE);
 
     if (prevNode != null && prevNode.getElementType() == PyElementTypes.IMPORT_ELEMENT &&
         result.getElementType() == PyElementTypes.IMPORT_ELEMENT) {
@@ -228,4 +231,22 @@ public class PyFromImportStatementImpl extends PyBaseElementImpl<PyFromImportSta
     }
     return ResolveImportUtil.resolveFromImportStatementSource(this, qName);
   }
+
+  @NotNull
+  @Override
+  public List<String> getFullyQualifiedObjectNames() {
+    final PyReferenceExpression source = getImportSource();
+    String prefix = "";
+    if (source != null) {
+      prefix = source.getText() + ".";
+    }
+    final List<String> unqualifiedNames = PyImportStatementImpl.getImportElementNames(getImportElements());
+
+    final List<String> result = new ArrayList<String>(unqualifiedNames.size());
+
+    for (final String unqualifiedName : unqualifiedNames) {
+      result.add(prefix + unqualifiedName);
+    }
+    return result;
+  }
 }

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

@@ -18,6 +18,7 @@ package com.jetbrains.python.psi.impl;
 import com.intellij.lang.ASTNode;
 import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.tree.TokenSet;
+import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.ArrayFactory;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.psi.PyElementVisitor;
@@ -26,6 +27,9 @@ import com.jetbrains.python.psi.PyImportStatement;
 import com.jetbrains.python.psi.stubs.PyImportStatementStub;
 import org.jetbrains.annotations.NotNull;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /**
  * @author yole
  */
@@ -66,4 +70,27 @@ public class PyImportStatementImpl extends PyBaseElementImpl<PyImportStatementSt
     PyPsiUtils.deleteAdjacentComma(this, child, getImportElements());
     super.deleteChildInternal(child);
   }
+
+  @NotNull
+  @Override
+  public List<String> getFullyQualifiedObjectNames() {
+    return getImportElementNames(getImportElements());
+  }
+
+  /**
+   * Returns list of qualified names of import elements filtering out nulls
+   * @param elements import elements
+   * @return list of qualified names
+   */
+  @NotNull
+  public static List<String> getImportElementNames(@NotNull final PyImportElement... elements) {
+    final List<String> result = new ArrayList<String>(elements.length);
+    for (final PyImportElement element : elements) {
+      final QualifiedName qName = element.getImportedQName();
+      if (qName != null) {
+        result.add(qName.toString());
+      }
+    }
+    return result;
+  }
 }