improve MROs for classes that inherit from generic protocols

Author: AlexWaygood

crates/red_knot_python_semantic/resources/mdtest/annotations/stdlib_typing_aliases.md

@@ -67,21 +67,24 @@ import typing
 
 ####################
 ### Built-ins
+####################
 
 class ListSubclass(typing.List): ...
 
-# revealed: tuple[Literal[ListSubclass], Literal[list], Literal[MutableSequence], Literal[Sequence], Literal[Reversible], Literal[Collection], Literal[Iterable], Literal[Container], @Todo(`Protocol[]` subscript), Literal[object]]
+# TODO: generic protocols
+# revealed: tuple[Literal[ListSubclass], Literal[list], Literal[MutableSequence], Literal[Sequence], Literal[Reversible], Literal[Collection], Literal[Iterable], Literal[Container], @Todo(`Protocol[]` subscript), @Todo(`Generic[]` subscript), Literal[object]]
 reveal_type(ListSubclass.__mro__)
 
 class DictSubclass(typing.Dict): ...
 
-# TODO
+# TODO: generic protocols
 # revealed: tuple[Literal[DictSubclass], Literal[dict], Literal[MutableMapping], Literal[Mapping], Literal[Collection], Literal[Iterable], Literal[Container], @Todo(`Protocol[]` subscript), @Todo(`Generic[]` subscript), Literal[object]]
 reveal_type(DictSubclass.__mro__)
 
 class SetSubclass(typing.Set): ...
 
-# revealed: tuple[Literal[SetSubclass], Literal[set], Literal[MutableSet], Literal[AbstractSet], Literal[Collection], Literal[Iterable], Literal[Container], @Todo(`Protocol[]` subscript), Literal[object]]
+# TODO: generic protocols
+# revealed: tuple[Literal[SetSubclass], Literal[set], Literal[MutableSet], Literal[AbstractSet], Literal[Collection], Literal[Iterable], Literal[Container], @Todo(`Protocol[]` subscript), @Todo(`Generic[]` subscript), Literal[object]]
 reveal_type(SetSubclass.__mro__)
 
 class FrozenSetSubclass(typing.FrozenSet): ...
@@ -92,10 +95,11 @@ reveal_type(FrozenSetSubclass.__mro__)
 
 ####################
 ### `collections`
+####################
 
 class ChainMapSubclass(typing.ChainMap): ...
 
-# TODO
+# TODO: generic protocols
 # revealed: tuple[Literal[ChainMapSubclass], Literal[ChainMap], Literal[MutableMapping], Literal[Mapping], Literal[Collection], Literal[Iterable], Literal[Container], @Todo(`Protocol[]` subscript), @Todo(`Generic[]` subscript), Literal[object]]
 reveal_type(ChainMapSubclass.__mro__)
 
@@ -113,7 +117,8 @@ reveal_type(DefaultDictSubclass.__mro__)
 
 class DequeSubclass(typing.Deque): ...
 
-# revealed: tuple[Literal[DequeSubclass], Literal[deque], Literal[MutableSequence], Literal[Sequence], Literal[Reversible], Literal[Collection], Literal[Iterable], Literal[Container], @Todo(`Protocol[]` subscript), Literal[object]]
+# TODO: generic protocols
+# revealed: tuple[Literal[DequeSubclass], Literal[deque], Literal[MutableSequence], Literal[Sequence], Literal[Reversible], Literal[Collection], Literal[Iterable], Literal[Container], @Todo(`Protocol[]` subscript), @Todo(`Generic[]` subscript), Literal[object]]
 reveal_type(DequeSubclass.__mro__)
 
 class OrderedDictSubclass(typing.OrderedDict): ...

crates/red_knot_python_semantic/resources/mdtest/subscript/tuple.md

@@ -117,6 +117,7 @@ from typing import Tuple
 
 class C(Tuple): ...
 
-# revealed: tuple[Literal[C], Literal[tuple], Literal[Sequence], Literal[Reversible], Literal[Collection], Literal[Iterable], Literal[Container], @Todo(`Protocol[]` subscript), Literal[object]]
+# TODO: generic protocols
+# revealed: tuple[Literal[C], Literal[tuple], Literal[Sequence], Literal[Reversible], Literal[Collection], Literal[Iterable], Literal[Container], @Todo(`Protocol[]` subscript), @Todo(`Generic[]` subscript), Literal[object]]
 reveal_type(C.__mro__)
 ```

crates/red_knot_python_semantic/src/types/class_base.rs

@@ -178,12 +178,20 @@ impl<'db> ClassBase<'db> {
     /// Iterate over the MRO of this base
     pub(super) fn mro(self, db: &'db dyn Db) -> impl Iterator<Item = ClassBase<'db>> {
         match self {
-            ClassBase::Dynamic(_) | ClassBase::Generic => {
-                Either::Left(Either::Left([self, ClassBase::object(db)].into_iter()))
-            }
-            ClassBase::Protocol => Either::Left(Either::Right(
+            ClassBase::Protocol => Either::Left(Either::Left(
                 [self, ClassBase::Generic, ClassBase::object(db)].into_iter(),
             )),
+            ClassBase::Dynamic(DynamicType::SubscriptedProtocol) => Either::Left(Either::Left(
+                [
+                    self,
+                    ClassBase::Dynamic(DynamicType::SubscriptedGeneric),
+                    ClassBase::object(db),
+                ]
+                .into_iter(),
+            )),
+            ClassBase::Dynamic(_) | ClassBase::Generic => {
+                Either::Left(Either::Right([self, ClassBase::object(db)].into_iter()))
+            }
             ClassBase::Class(class) => Either::Right(class.iter_mro(db)),
         }
     }

crates/red_knot_python_semantic/src/types/infer.rs

@@ -6358,15 +6358,9 @@ impl<'db> TypeInferenceBuilder<'db> {
 
                     // TODO: properly handle old-style generics; get rid of this temporary hack
                     if !value_ty.into_class_literal().is_some_and(|class| {
-                        class.iter_mro(self.db(), None).any(|base| {
-                            matches!(
-                                base,
-                                ClassBase::Dynamic(
-                                    DynamicType::SubscriptedGeneric
-                                        | DynamicType::SubscriptedProtocol,
-                                )
-                            )
-                        })
+                        class
+                            .iter_mro(self.db(), None)
+                            .any(|base| base == ClassBase::Dynamic(DynamicType::SubscriptedGeneric))
                     }) {
                         report_non_subscriptable(
                             &self.context,