Support mock.patch members in Python 2 (PY-23632)

python/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibClassMembersProvider.java

@@ -3,6 +3,7 @@ package com.jetbrains.python.codeInsight.stdlib;
 
 import com.intellij.openapi.util.Key;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.codeInsight.PyCustomMember;
 import com.jetbrains.python.psi.*;
@@ -24,9 +25,6 @@ public class PyStdlibClassMembersProvider extends PyClassMembersProviderBase {
   @NotNull
   private static final Key<List<PyCustomMember>> SOCKET_MEMBERS_KEY = Key.create("socket.members");
 
-  @NotNull
-  public static final List<PyCustomMember> MOCK_PATCH_MEMBERS = calcMockPatchMembers();
-
   @NotNull
   @Override
   public Collection<PyCustomMember> getMembers(PyClassType classType, PsiElement location, @NotNull TypeEvalContext context) {
@@ -45,7 +43,7 @@ public class PyStdlibClassMembersProvider extends PyClassMembersProviderBase {
     if (location instanceof PyReferenceExpression) {
       final PyExpression qualifier = ((PyReferenceExpression)location).getQualifier();
       if (qualifier instanceof PyReferenceExpression && referenceToMockPatch((PyReferenceExpression)qualifier, context)) {
-        return MOCK_PATCH_MEMBERS;
+        return calcMockPatchMembers(qualifier);
       }
     }
 
@@ -60,7 +58,7 @@ public class PyStdlibClassMembersProvider extends PyClassMembersProviderBase {
                                   @NotNull PyResolveContext resolveContext) {
     final TypeEvalContext context = resolveContext.getTypeEvalContext();
     if (location instanceof PyReferenceExpression && referenceToMockPatch((PyReferenceExpression)location, context)) {
-      for (PyCustomMember member : MOCK_PATCH_MEMBERS) {
+      for (PyCustomMember member : calcMockPatchMembers(location)) {
         if (name.equals(member.getName())) {
           return member.resolve(location, resolveContext);
         }
@@ -73,7 +71,8 @@ public class PyStdlibClassMembersProvider extends PyClassMembersProviderBase {
   public static boolean referenceToMockPatch(@NotNull PyReferenceExpression referenceExpression, @NotNull TypeEvalContext context) {
     final PyType type = context.getType(referenceExpression);
     if (type instanceof PyFunctionType) {
-      return "unittest.mock.patch".equals(((PyFunctionType)type).getCallable().getQualifiedName());
+      final String callableQName = ((PyFunctionType)type).getCallable().getQualifiedName();
+      return mockPatchQName(referenceExpression).toString().equals(callableQName);
     }
     return false;
   }
@@ -86,13 +85,20 @@ public class PyStdlibClassMembersProvider extends PyClassMembersProviderBase {
   }
 
   @NotNull
-  private static List<PyCustomMember> calcMockPatchMembers() {
-    final String[] members = new String[]{"object", "dict", "multiple", "stopall", "TEST_PREFIX"};
-    final String moduleQName = "unittest.mock";
+  public static List<PyCustomMember> calcMockPatchMembers(@NotNull PsiElement anchor) {
+    final String[] members = {"object", "dict", "multiple", "stopall", "TEST_PREFIX"};
+    final String moduleQName = mockPatchQName(anchor).removeLastComponent().toString();
 
     return ContainerUtil.map(members, member -> new PyCustomMember(member).resolvesTo(moduleQName).toAssignment("patch." + member));
   }
 
+  @NotNull
+  private static QualifiedName mockPatchQName(@NotNull PsiElement anchor) {
+    return LanguageLevel.forElement(anchor).isPython2()
+           ? QualifiedName.fromComponents("mock", "mock", "patch")
+           : QualifiedName.fromComponents("unittest", "mock", "patch");
+  }
+
   private static void addMethodsFromAttr(PyFile socketFile, List<PyCustomMember> result, final String attrName) {
     final PyTargetExpression socketMethods = socketFile.findTopLevelAttribute(attrName);
     if (socketMethods != null) {

python/src/com/jetbrains/python/inspections/stdlib/PyStdlibInspectionExtension.kt

@@ -38,7 +38,7 @@ class PyStdlibInspectionExtension : PyInspectionExtension() {
       val qualifier = node.qualifier
       if (qualifier is PyReferenceExpression) {
         return PyStdlibClassMembersProvider.referenceToMockPatch(qualifier, context) &&
-               PyStdlibClassMembersProvider.MOCK_PATCH_MEMBERS.find { it.name == node.name } != null
+               PyStdlibClassMembersProvider.calcMockPatchMembers(qualifier).any { it.name == node.name }
       }
     }
 

python/testData/completion/mockPatchObject1Py2/a.after.py

@@ -0,0 +1,3 @@
+from mock import patch
+
+patch.object

python/testData/completion/mockPatchObject1Py2/a.py

@@ -0,0 +1,3 @@
+from mock import patch
+
+patch.obj<caret>

python/testData/completion/mockPatchObject1Py2/mock/__init__.py

@@ -0,0 +1 @@
+from mock.mock import *

python/testData/completion/mockPatchObject1Py2/mock/mock.py

@@ -0,0 +1,4 @@
+def patch():
+    pass
+
+patch.object = None

python/testData/completion/mockPatchObject2Py2/a.after.py

@@ -0,0 +1,3 @@
+import mock
+
+mock.patch.object

python/testData/completion/mockPatchObject2Py2/a.py

@@ -0,0 +1,3 @@
+import mock
+
+mock.patch.obj<caret>

python/testData/completion/mockPatchObject2Py2/mock/__init__.py

@@ -0,0 +1 @@
+from mock.mock import *

python/testData/completion/mockPatchObject2Py2/mock/mock.py

@@ -0,0 +1,4 @@
+def patch():
+    pass
+
+patch.object = None

python/testData/inspections/PyUnresolvedReferencesInspection/MockPatchObject/a.py

@@ -0,0 +1,7 @@
+from mock import patch
+
+print(patch.object)
+
+import mock
+
+print(mock.patch.object)

python/testData/inspections/PyUnresolvedReferencesInspection/MockPatchObject/mock/__init__.py

@@ -0,0 +1 @@
+from mock.mock import *

python/testData/inspections/PyUnresolvedReferencesInspection/MockPatchObject/mock/mock.py

@@ -0,0 +1,4 @@
+def patch():
+    pass
+
+patch.object = None

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

@@ -1235,6 +1235,16 @@ public class PythonCompletionTest extends PyTestCase {
     );
   }
 
+  // PY-23632
+  public void testMockPatchObject1Py2() {
+    doMultiFileTest();
+  }
+
+  // PY-23632
+  public void testMockPatchObject2Py2() {
+    doMultiFileTest();
+  }
+
   @Override
   protected String getTestDataPath() {
     return super.getTestDataPath() + "/completion";

python/testSrc/com/jetbrains/python/inspections/PyUnresolvedReferencesInspectionTest.java

@@ -681,6 +681,11 @@ public class PyUnresolvedReferencesInspectionTest extends PyInspectionTestCase {
     doMultiFileTest();
   }
 
+  // PY-23632
+  public void testMockPatchObject() {
+    doMultiFileTest();
+  }
+
   @NotNull
   @Override
   protected Class<? extends PyInspection> getInspectionClass() {