From 4495d85bb49659cf0a246b0f91ee2118f1ab91d9 Mon Sep 17 00:00:00 2001
From: Ekaterina Tuzova <Ekaterina.Tuzova@jetbrains.com>
Date: Tue, 5 May 2015 14:44:22 +0300
Subject: [PATCH] PY-15366 No completion for matplotlib figure

matplotlib uses numpy docstrings. Added numpy to the user code ( due to PY-4717 Support for NumpyDoc)
---
 .../NumpyDocStringTypeProvider.java           | 25 +++++++++++++++----
 .../python/documentation/DocStringFormat.java |  3 ++-
 .../PyDocumentationSettings.java              |  4 +++
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/python/src/com/jetbrains/numpy/codeInsight/NumpyDocStringTypeProvider.java b/python/src/com/jetbrains/numpy/codeInsight/NumpyDocStringTypeProvider.java
index bee35b9ec768..389e7c38f7d0 100644
--- a/python/src/com/jetbrains/numpy/codeInsight/NumpyDocStringTypeProvider.java
+++ b/python/src/com/jetbrains/numpy/codeInsight/NumpyDocStringTypeProvider.java
@@ -17,6 +17,8 @@ package com.jetbrains.numpy.codeInsight;
 
 import com.google.common.collect.Lists;
 import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.module.Module;
+import com.intellij.openapi.module.ModuleUtilCore;
 import com.intellij.openapi.util.Ref;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.openapi.vfs.VirtualFile;
@@ -24,6 +26,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.jetbrains.numpy.documentation.NumPyDocString;
 import com.jetbrains.numpy.documentation.NumPyDocStringParameter;
+import com.jetbrains.python.documentation.PyDocumentationSettings;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyBuiltinCache;
 import com.jetbrains.python.psi.impl.PyExpressionCodeFragmentImpl;
@@ -83,7 +86,7 @@ public class NumpyDocStringTypeProvider extends PyTypeProviderBase {
   @Nullable
   @Override
   public PyType getCallType(@NotNull PyFunction function, @Nullable PyCallSiteExpression callSite, @NotNull TypeEvalContext context) {
-    if (isInsideNumPy(function)) {
+    if (isApplicable(function)) {
       final PyExpression callee = callSite instanceof PyCallExpression ? ((PyCallExpression)callSite).getCallee() : null;
       final NumPyDocString docString = NumPyDocString.forFunction(function, callee);
       if (docString != null) {
@@ -140,7 +143,7 @@ public class NumpyDocStringTypeProvider extends PyTypeProviderBase {
   @Nullable
   @Override
   public Ref<PyType> getParameterType(@NotNull PyNamedParameter parameter, @NotNull PyFunction function, @NotNull TypeEvalContext context) {
-    if (isInsideNumPy(function)) {
+    if (isApplicable(function)) {
       final String name = parameter.getName();
       if (name != null) {
         final PyType type = getParameterType(function, name);
@@ -153,19 +156,31 @@ public class NumpyDocStringTypeProvider extends PyTypeProviderBase {
   }
 
   private static boolean isInsideNumPy(@NotNull PsiElement element) {
-    if (ApplicationManager.getApplication().isUnitTestMode()) return true;
     final PsiFile file = element.getContainingFile();
+
     if (file != null) {
       final PyPsiFacade facade = getPsiFacade(element);
       final VirtualFile virtualFile = file.getVirtualFile();
       if (virtualFile != null) {
         final String name = facade.findShortestImportableName(virtualFile, element);
-        return name != null && name.startsWith("numpy.");
+        return name != null && (name.startsWith("numpy.") || name.startsWith("matplotlib."));
       }
     }
     return false;
   }
 
+  private static boolean isApplicable(@NotNull PsiElement element) {
+    if (ApplicationManager.getApplication().isUnitTestMode()) return true;
+    final Module module = ModuleUtilCore.findModuleForPsiElement(element);
+    if (module != null){
+      if (PyDocumentationSettings.getInstance(module).isNumpyFormat(element.getContainingFile())) {
+        return true;
+      }
+    }
+
+    return isInsideNumPy(element);
+  }
+
   private static PyPsiFacade getPsiFacade(@NotNull PsiElement anchor) {
     return PyPsiFacade.getInstance(anchor.getProject());
   }
@@ -237,7 +252,7 @@ public class NumpyDocStringTypeProvider extends PyTypeProviderBase {
 
   private static boolean isUfuncType(@NotNull PsiElement anchor, @NotNull final String typeString) {
     for (String typeName : NumPyDocString.getNumpyUnionType(typeString)) {
-      if (anchor instanceof PyFunction && NumpyUfuncs.isUFunc(((PyFunction)anchor).getName()) &&
+      if (anchor instanceof PyFunction && isInsideNumPy(anchor) && NumpyUfuncs.isUFunc(((PyFunction)anchor).getName()) &&
           ("array_like".equals(typeName) || "ndarray".equals(typeName))) {
         return true;
       }
diff --git a/python/src/com/jetbrains/python/documentation/DocStringFormat.java b/python/src/com/jetbrains/python/documentation/DocStringFormat.java
index 094e47d5554a..164917820b06 100644
--- a/python/src/com/jetbrains/python/documentation/DocStringFormat.java
+++ b/python/src/com/jetbrains/python/documentation/DocStringFormat.java
@@ -26,8 +26,9 @@ public class DocStringFormat {
   public static final String PLAIN = "Plain";
   public static final String EPYTEXT = "Epytext";
   public static final String REST = "reStructuredText";
+  public static final String NUMPY = "NumPy";
 
-  public static final List<String> ALL = ImmutableList.of(PLAIN, EPYTEXT, REST);
+  public static final List<String> ALL = ImmutableList.of(PLAIN, EPYTEXT, REST, NUMPY);
 
   private DocStringFormat() {
   }
diff --git a/python/src/com/jetbrains/python/documentation/PyDocumentationSettings.java b/python/src/com/jetbrains/python/documentation/PyDocumentationSettings.java
index fe36834f6c5b..f17e8ff3f834 100644
--- a/python/src/com/jetbrains/python/documentation/PyDocumentationSettings.java
+++ b/python/src/com/jetbrains/python/documentation/PyDocumentationSettings.java
@@ -53,6 +53,10 @@ public class PyDocumentationSettings implements PersistentStateComponent<PyDocum
     return isFormat(file, DocStringFormat.REST);
   }
 
+  public boolean isNumpyFormat(PsiFile file) {
+    return isFormat(file, DocStringFormat.NUMPY);
+  }
+
   public boolean isPlain(PsiFile file) {
     return isFormat(file, DocStringFormat.PLAIN);
   }