From 71b52b83e4c6f4705fecfa556324ae2cfd8f2ec5 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Mon, 7 Oct 2024 19:43:47 +0100
Subject: [PATCH] [red-knot] Allow `type[]` to be subscripted (#13667)

Fixed a TODO by adding another TODO. It's the red-knot way!

## Summary

`builtins.type` can be subscripted at runtime on Python 3.9+, even
though it has no `__class_getitem__` method and its metaclass (which
is... itself) has no `__getitem__` method. The special case is
[hardcoded directly into `PyObject_GetItem` in
CPython](https://github.com/python/cpython/blob/744caa8ef42ab67c6aa20cd691e078721e72e22a/Objects/abstract.c#L181-L184).
We just have to replicate the special case in our semantic model.

This will fail at runtime on Python <3.9. However, there's a bunch of
outstanding questions (detailed in the TODO comment I added) regarding
how we deal with subscriptions of other generic types on lower Python
versions. Since we want to avoid too many false positives for now, I
haven't tried to address this; I've just made `type` subscriptable on
all Python versions.

## Test Plan

`cargo test -p red_knot_python_semantic --lib`
---
 crates/red_knot_python_semantic/src/types.rs  |  9 +++++++--
 .../src/types/infer.rs                        | 19 ++++++++++++-------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
index ce84a55e6d..3e686b1636 100644
--- a/crates/red_knot_python_semantic/src/types.rs
+++ b/crates/red_knot_python_semantic/src/types.rs
@@ -837,6 +837,7 @@ pub enum KnownClass {
     Set,
     Dict,
     // Types
+    GenericAlias,
     ModuleType,
     FunctionType,
     // Typeshed
@@ -857,6 +858,7 @@ impl<'db> KnownClass {
             Self::Dict => "dict",
             Self::List => "list",
             Self::Type => "type",
+            Self::GenericAlias => "GenericAlias",
             Self::ModuleType => "ModuleType",
             Self::FunctionType => "FunctionType",
             Self::NoneType => "NoneType",
@@ -880,7 +882,9 @@ impl<'db> KnownClass {
             | Self::Tuple
             | Self::Set
             | Self::Dict => builtins_symbol_ty(db, self.as_str()),
-            Self::ModuleType | Self::FunctionType => types_symbol_ty(db, self.as_str()),
+            Self::GenericAlias | Self::ModuleType | Self::FunctionType => {
+                types_symbol_ty(db, self.as_str())
+            }
             Self::NoneType => typeshed_symbol_ty(db, self.as_str()),
         }
     }
@@ -910,6 +914,7 @@ impl<'db> KnownClass {
             "set" => Some(Self::Set),
             "dict" => Some(Self::Dict),
             "list" => Some(Self::List),
+            "GenericAlias" => Some(Self::GenericAlias),
             "NoneType" => Some(Self::NoneType),
             "ModuleType" => Some(Self::ModuleType),
             "FunctionType" => Some(Self::FunctionType),
@@ -934,7 +939,7 @@ impl<'db> KnownClass {
             | Self::Tuple
             | Self::Set
             | Self::Dict => module.name() == "builtins",
-            Self::ModuleType | Self::FunctionType => module.name() == "types",
+            Self::GenericAlias | Self::ModuleType | Self::FunctionType => module.name() == "types",
             Self::NoneType => matches!(module.name().as_str(), "_typeshed" | "types"),
         }
     }
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index 349ef57651..f34892ff5d 100644
--- a/crates/red_knot_python_semantic/src/types/infer.rs
+++ b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -2828,6 +2828,13 @@ impl<'db> TypeInferenceBuilder<'db> {
 
                 // Otherwise, if the value is itself a class and defines `__class_getitem__`,
                 // return its return type.
+                //
+                // TODO: lots of classes are only subscriptable at runtime on Python 3.9+,
+                // *but* we should also allow them to be subscripted in stubs
+                // (and in annotations if `from __future__ import annotations` is enabled),
+                // even if the target version is Python 3.8 or lower,
+                // despite the fact that there will be no corresponding `__class_getitem__`
+                // method in these `sys.version_info` branches.
                 if value_ty.is_class(self.db) {
                     let dunder_class_getitem_method = value_ty.member(self.db, "__class_getitem__");
                     if !dunder_class_getitem_method.is_unbound() {
@@ -2848,6 +2855,11 @@ impl<'db> TypeInferenceBuilder<'db> {
                             });
                     }
 
+                    if matches!(value_ty, Type::Class(class) if class.is_known(self.db, KnownClass::Type))
+                    {
+                        return KnownClass::GenericAlias.to_instance(self.db);
+                    }
+
                     self.non_subscriptable_diagnostic(
                         (&**value).into(),
                         value_ty,
@@ -6194,13 +6206,6 @@ mod tests {
         )?;
 
         let expected_diagnostics = &[
-            // TODO: these `__class_getitem__` diagnostics are all false positives:
-            // (`builtins.type` is unique at runtime
-            //  as it can be subscripted even though it has no `__class_getitem__` method)
-            "Cannot subscript object of type `Literal[type]` with no `__class_getitem__` method",
-            "Cannot subscript object of type `Literal[type]` with no `__class_getitem__` method",
-            "Cannot subscript object of type `Literal[type]` with no `__class_getitem__` method",
-            "Cannot subscript object of type `Literal[type]` with no `__class_getitem__` method",
             // Should be `AttributeError`:
             "Revealed type is `@Todo`",
             // Should be `OSError | RuntimeError`: