[ty] Fix protocol interface inference for stub protocols and subprotocols (#19950)

crates/ty_python_semantic/resources/mdtest/protocols.md

@@ -397,7 +397,7 @@ To see the kinds and types of the protocol members, you can use the debugging ai
 
 ```py
 from ty_extensions import reveal_protocol_interface
-from typing import SupportsIndex, SupportsAbs
+from typing import SupportsIndex, SupportsAbs, ClassVar
 
 # error: [revealed-type] "Revealed protocol interface: `{"method_member": MethodMember(`(self) -> bytes`), "x": AttributeMember(`int`), "y": PropertyMember { getter: `def y(self) -> str` }, "z": PropertyMember { getter: `def z(self) -> int`, setter: `def z(self, z: int) -> None` }}`"
 reveal_protocol_interface(Foo)
@@ -415,6 +415,33 @@ reveal_protocol_interface("foo")
 #
 # error: [invalid-argument-type] "Invalid argument to `reveal_protocol_interface`: Only protocol classes can be passed to `reveal_protocol_interface`"
 reveal_protocol_interface(SupportsAbs[int])
+
+class BaseProto(Protocol):
+    def member(self) -> int: ...
+
+class SubProto(BaseProto, Protocol):
+    def member(self) -> bool: ...
+
+# error: [revealed-type] "Revealed protocol interface: `{"member": MethodMember(`(self) -> int`)}`"
+reveal_protocol_interface(BaseProto)
+
+# error: [revealed-type] "Revealed protocol interface: `{"member": MethodMember(`(self) -> bool`)}`"
+reveal_protocol_interface(SubProto)
+
+class ProtoWithClassVar(Protocol):
+    x: ClassVar[int]
+
+# error: [revealed-type] "Revealed protocol interface: `{"x": AttributeMember(`int`; ClassVar)}`"
+reveal_protocol_interface(ProtoWithClassVar)
+
+class ProtocolWithDefault(Protocol):
+    x: int = 0
+
+# We used to incorrectly report this as having an `x: Literal[0]` member;
+# declared types should take priority over inferred types for protocol interfaces!
+#
+# error: [revealed-type] "Revealed protocol interface: `{"x": AttributeMember(`int`)}`"
+reveal_protocol_interface(ProtocolWithDefault)
 ```
 
 Certain special attributes and methods are not considered protocol members at runtime, and should
@@ -623,13 +650,26 @@ class HasXWithDefault(Protocol):
 class FooWithZero:
     x: int = 0
 
-# TODO: these should pass
+static_assert(is_subtype_of(FooWithZero, HasXWithDefault))
-static_assert(is_subtype_of(FooWithZero, HasXWithDefault))  # error: [static-assert-error]
+static_assert(is_assignable_to(FooWithZero, HasXWithDefault))
-static_assert(is_assignable_to(FooWithZero, HasXWithDefault))  # error: [static-assert-error]
+
-static_assert(not is_subtype_of(Foo, HasXWithDefault))
+# TODO: whether or not any of these four assertions should pass is not clearly specified.
-static_assert(not is_assignable_to(Foo, HasXWithDefault))
+#
-static_assert(not is_subtype_of(Qux, HasXWithDefault))
+# A test in the typing conformance suite implies that they all should:
-static_assert(not is_assignable_to(Qux, HasXWithDefault))
+# that a nominal class with an instance attribute `x`
+# (*without* a default value on the class body)
+# should be understood as satisfying a protocol that has an attribute member `x`
+# even if the protocol's `x` member has a default value on the class body.
+#
+# See <https://github.com/python/typing/blob/d4f39b27a4a47aac8b6d4019e1b0b5b3156fabdc/conformance/tests/protocols_definition.py#L56-L79>.
+#
+# The implications of this for meta-protocols are not clearly spelled out, however,
+# and the fact that attribute members on protocols can have defaults is only mentioned
+# in a throwaway comment in the spec's prose.
+static_assert(is_subtype_of(Foo, HasXWithDefault))
+static_assert(is_assignable_to(Foo, HasXWithDefault))
+static_assert(is_subtype_of(Qux, HasXWithDefault))
+static_assert(is_assignable_to(Qux, HasXWithDefault))
 
 class HasClassVarX(Protocol):
     x: ClassVar[int]
@@ -2127,6 +2167,30 @@ static_assert(is_subtype_of(type[Baz], type[Foo]))  # error: [static-assert-erro
 static_assert(is_subtype_of(TypeOf[Baz], type[Foo]))  # error: [static-assert-error]
 ```
 
+## Regression test for `ClassVar` members in stubs
+
+In an early version of our protocol implementation, we didn't retain the `ClassVar` qualifier for
+protocols defined in stub files.
+
+`stub.pyi`:
+
+```pyi
+from typing import ClassVar, Protocol
+
+class Foo(Protocol):
+    x: ClassVar[int]
+```
+
+`main.py`:
+
+```py
+from stub import Foo
+from ty_extensions import reveal_protocol_interface
+
+# error: [revealed-type] "Revealed protocol interface: `{"x": AttributeMember(`int`; ClassVar)}`"
+reveal_protocol_interface(Foo)
+```
+
 ## TODO
 
 Add tests for:

crates/ty_python_semantic/src/types/protocol_class.rs

@@ -4,6 +4,7 @@ use std::{collections::BTreeMap, ops::Deref};
 use itertools::Itertools;
 
 use ruff_python_ast::name::Name;
+use rustc_hash::FxHashMap;
 
 use super::TypeVarVariance;
 use crate::semantic_index::place_table;
@@ -286,6 +287,7 @@ impl<'db> ProtocolMemberData<'db> {
         struct ProtocolMemberDataDisplay<'db> {
             db: &'db dyn Db,
             data: ProtocolMemberKind<'db>,
+            qualifiers: TypeQualifiers,
         }
 
         impl std::fmt::Display for ProtocolMemberDataDisplay<'_> {
@@ -305,7 +307,12 @@ impl<'db> ProtocolMemberData<'db> {
                         d.finish()
                     }
                     ProtocolMemberKind::Other(ty) => {
-                        write!(f, "AttributeMember(`{}`)", ty.display(self.db))
+                        f.write_str("AttributeMember(")?;
+                        write!(f, "`{}`", ty.display(self.db))?;
+                        if self.qualifiers.contains(TypeQualifiers::CLASS_VAR) {
+                            f.write_str("; ClassVar")?;
+                        }
+                        f.write_char(')')
                     }
                 }
             }
@@ -314,6 +321,7 @@ impl<'db> ProtocolMemberData<'db> {
         ProtocolMemberDataDisplay {
             db,
             data: self.kind,
+            qualifiers: self.qualifiers,
         }
     }
 }
@@ -517,6 +525,12 @@ enum BoundOnClass {
     No,
 }
 
+impl BoundOnClass {
+    const fn is_yes(self) -> bool {
+        matches!(self, BoundOnClass::Yes)
+    }
+}
+
 /// Inner Salsa query for [`ProtocolClassLiteral::interface`].
 #[salsa::tracked(cycle_fn=proto_interface_cycle_recover, cycle_initial=proto_interface_cycle_initial, heap_size=ruff_memory_usage::heap_size)]
 fn cached_protocol_interface<'db>(
@@ -533,22 +547,8 @@ fn cached_protocol_interface<'db>(
         let parent_scope = parent_protocol.body_scope(db);
         let use_def_map = use_def_map(db, parent_scope);
         let place_table = place_table(db, parent_scope);
+        let mut direct_members = FxHashMap::default();
 
-        members.extend(
-            use_def_map
-                .all_end_of_scope_symbol_declarations()
-                .map(|(symbol_id, declarations)| {
-                    (
-                        symbol_id,
-                        place_from_declarations(db, declarations).ignore_conflicting_declarations(),
-                    )
-                })
-                .filter_map(|(symbol_id, place)| {
-                    place
-                        .place
-                        .ignore_possibly_unbound()
-                        .map(|ty| (symbol_id, ty, place.qualifiers, BoundOnClass::No))
-                })
         // Bindings in the class body that are not declared in the class body
         // are not valid protocol members, and we plan to emit diagnostics for them
         // elsewhere. Invalid or not, however, it's important that we still consider
@@ -557,46 +557,60 @@ fn cached_protocol_interface<'db>(
         // members at runtime, and it's important that we accurately understand
         // type narrowing that uses `isinstance()` or `issubclass()` with
         // runtime-checkable protocols.
-                .chain(use_def_map.all_end_of_scope_symbol_bindings().filter_map(
+        for (symbol_id, bindings) in use_def_map.all_end_of_scope_symbol_bindings() {
-                    |(symbol_id, bindings)| {
+            let Some(ty) = place_from_bindings(db, bindings).ignore_possibly_unbound() else {
-                        place_from_bindings(db, bindings)
+                continue;
-                            .ignore_possibly_unbound()
+            };
-                            .map(|ty| (symbol_id, ty, TypeQualifiers::default(), BoundOnClass::Yes))
+            direct_members.insert(
-                    },
+                symbol_id,
-                ))
+                (ty, TypeQualifiers::default(), BoundOnClass::Yes),
-                .map(|(symbol_id, member, qualifiers, bound_on_class)| {
+            );
-                    (
-                        place_table.symbol(symbol_id).name(),
-                        member,
-                        qualifiers,
-                        bound_on_class,
-                    )
-                })
-                .filter(|(name, _, _, _)| !excluded_from_proto_members(name))
-                .map(|(name, ty, qualifiers, bound_on_class)| {
-                    let kind = match (ty, bound_on_class) {
-                        // TODO: if the getter or setter is a function literal, we should
-                        // upcast it to a `CallableType` so that two protocols with identical property
-                        // members are recognized as equivalent.
-                        (Type::PropertyInstance(property), _) => {
-                            ProtocolMemberKind::Property(property)
         }
-                        (Type::Callable(callable), BoundOnClass::Yes)
+
-                            if callable.is_function_like(db) =>
+        for (symbol_id, declarations) in use_def_map.all_end_of_scope_symbol_declarations() {
+            let place = place_from_declarations(db, declarations).ignore_conflicting_declarations();
+            if let Some(new_type) = place.place.ignore_possibly_unbound() {
+                direct_members
+                    .entry(symbol_id)
+                    .and_modify(|(ty, quals, _)| {
+                        *ty = new_type;
+                        *quals = place.qualifiers;
+                    })
+                    .or_insert((new_type, place.qualifiers, BoundOnClass::No));
+            }
+        }
+
+        for (symbol_id, (ty, qualifiers, bound_on_class)) in direct_members {
+            let name = place_table.symbol(symbol_id).name();
+            if excluded_from_proto_members(name) {
+                continue;
+            }
+            if members.contains_key(name) {
+                continue;
+            }
+
+            let member = match ty {
+                Type::PropertyInstance(property) => ProtocolMemberKind::Property(property),
+                Type::Callable(callable)
+                    if bound_on_class.is_yes() && callable.is_function_like(db) =>
                 {
                     ProtocolMemberKind::Method(callable)
                 }
-                        (Type::FunctionLiteral(function), BoundOnClass::Yes) => {
+                Type::FunctionLiteral(function) if bound_on_class.is_yes() => {
                     ProtocolMemberKind::Method(function.into_callable_type(db))
                 }
                 _ => ProtocolMemberKind::Other(ty),
             };
 
-                    let member = ProtocolMemberData { kind, qualifiers };
+            members.insert(
-                    (name.clone(), member)
+                name.clone(),
-                }),
+                ProtocolMemberData {
+                    kind: member,
+                    qualifiers,
+                },
             );
         }
+    }
 
     ProtocolInterface::new(db, members)
 }