PY-36062: Add resolve attributes matching PyModuleType.MODULE_MEMBERS to types.ModuleType attributes. Add matching PyClassType: ModuleType to PyModuleType.

GitOrigin-RevId: d2988397e4a58f9170a90cd6a3251054dedf896f

python/python-common-tests/com/jetbrains/python/PyCommonResolveTest.java

@@ -1532,4 +1532,12 @@ public abstract class PyCommonResolveTest extends PyCommonResolveTestCase {
     final TypeEvalContext context = TypeEvalContext.codeInsightFallback(myFixture.getProject());
     assertEmpty(file.findTopLevelAttribute("t").multiResolveAssignedValue(PyResolveContext.defaultContext().withTypeEvalContext(context)));
   }
+
+  // PY-36062
+  public void testModuleTypeAttributes() {
+    myFixture.copyDirectoryToProject("resolve/" + getTestName(false), "");
+    final PyTargetExpression target = assertResolvesTo(PyTargetExpression.class, "__name__");
+    assertEquals("ModuleType", target.getContainingClass().getName());
+    assertEquals("types.pyi", target.getContainingFile().getName());
+  }
 }

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

@@ -42,7 +42,7 @@ import static com.jetbrains.python.psi.PyUtil.inSameFile;
 public class PyModuleType implements PyType { // Modules don't descend from object
   @NotNull private final PyFile myModule;
 
-  public static final ImmutableSet<String> MODULE_MEMBERS = ImmutableSet.of(
+  private static final ImmutableSet<String> MODULE_MEMBERS = ImmutableSet.of(
     "__name__", "__file__", "__path__", "__doc__", "__dict__", "__package__");
 
   public PyModuleType(@NotNull PyFile source) {
@@ -55,12 +55,23 @@ public class PyModuleType implements PyType { // Modules don't descend from obje
     return myModule;
   }
 
+  @Nullable
+  public PyClassType getModuleClassType() {
+    return PyClassTypeImpl.createTypeByQName(myModule, "types.ModuleType", false);
+  }
+
   @Nullable
   @Override
   public List<? extends RatedResolveResult> resolveMember(@NotNull final String name,
                                                           @Nullable PyExpression location,
                                                           @NotNull AccessDirection direction,
                                                           @NotNull PyResolveContext resolveContext) {
+    if (MODULE_MEMBERS.contains(name)) {
+      var type = getModuleClassType();
+      if (type != null) {
+        return type.resolveMember(name, location, direction, resolveContext);
+      }
+    }
     return resolveMemberInPackageOrModule(null, myModule, name, location, resolveContext);
   }
 
@@ -562,6 +573,6 @@ public class PyModuleType implements PyType { // Modules don't descend from obje
 
   @Override
   public @Nullable PyQualifiedNameOwner getDeclarationElement() {
-    return PyPsiFacade.getInstance(myModule.getProject()).createClassByQName("types.ModuleType", myModule);
+    return ObjectUtils.doIfNotNull(getModuleClassType(), PyClassType::getPyClass);
   }
 }

python/python-psi-impl/src/com/jetbrains/python/psi/types/PyTypeChecker.java

@@ -156,6 +156,10 @@ public final class PyTypeChecker {
       return Optional.of(actual instanceof PyModuleType && ((PyModuleType)expected).getModule() == ((PyModuleType)actual).getModule());
     }
 
+    if (expected instanceof PyClassType && actual instanceof PyModuleType) {
+      return match(expected, ((PyModuleType)actual).getModuleClassType(), context);
+    }
+
     return Optional.of(matchNumericTypes(expected, actual));
   }
 

python/testData/inspections/PyTypeCheckerInspection/ModuleTypeParameter/a.py

@@ -0,0 +1,11 @@
+import module
+from types import ModuleType
+
+def foo(m: ModuleType):
+    pass
+
+def bar(m):
+    return m.__name__
+
+foo(module)
+bar(module)

python/testData/resolve/ModuleTypeAttributes.py

@@ -0,0 +1,3 @@
+import b
+b.__name__
+#    <ref>

python/testData/resolve/ModuleTypeAttributes/b.py

@@ -0,0 +1 @@
+__name__ = "abc"

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

@@ -494,4 +494,10 @@ public class Py3TypeCheckerInspectionTest extends PyInspectionTestCase {
                          "mymap(myfoo, [1, 2, 3])\n")
     );
   }
+
+  // PY-36062
+  public void testModuleTypeParameter() {
+    // `types.ModuleType` class qualified name is `_importlib_modulespec.ModuleType` in Python 3
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doMultiFileTest);
+  }
 }

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

@@ -1269,4 +1269,9 @@ public class PyTypeCheckerInspectionTest extends PyInspectionTestCase {
   public void testClassInstanceAgainstProtocolDunderCall() {
     runWithLanguageLevel(LanguageLevel.getLatest(), this::doTest);
   }
+
+  // PY-36062
+  public void testModuleTypeParameter() {
+    runWithLanguageLevel(LanguageLevel.getLatest(), this::doMultiFileTest);
+  }
 }