[ty] Strict validation of protocol members (#17750)

crates/ty_python_semantic/src/semantic_index/definition.rs

@@ -702,6 +702,10 @@ impl DefinitionKind<'_> {
         )
     }
 
+    pub(crate) const fn is_unannotated_assignment(&self) -> bool {
+        matches!(self, DefinitionKind::Assignment(_))
+    }
+
     pub(crate) fn as_typevar(&self) -> Option<&AstNodeRef<ast::TypeParamTypeVar>> {
         match self {
             DefinitionKind::TypeVar(type_var) => Some(type_var),

crates/ty_python_semantic/src/types/diagnostic.rs

@@ -7,8 +7,9 @@ use super::{
 };
 use crate::lint::{Level, LintRegistryBuilder, LintStatus};
 use crate::semantic_index::SemanticIndex;
+use crate::semantic_index::definition::Definition;
+use crate::semantic_index::place::{PlaceTable, ScopedPlaceId};
 use crate::suppression::FileSuppressionId;
-use crate::types::LintDiagnosticGuard;
 use crate::types::class::{Field, SolidBase, SolidBaseKind};
 use crate::types::function::KnownFunction;
 use crate::types::string_annotation::{
@@ -16,6 +17,9 @@ use crate::types::string_annotation::{
     IMPLICIT_CONCATENATED_STRING_TYPE_ANNOTATION, INVALID_SYNTAX_IN_FORWARD_ANNOTATION,
     RAW_STRING_TYPE_ANNOTATION,
 };
+use crate::types::{
+    DynamicType, LintDiagnosticGuard, Protocol, ProtocolInstanceType, SubclassOfInner, binding_type,
+};
 use crate::types::{SpecialFormType, Type, protocol_class::ProtocolClassLiteral};
 use crate::util::diagnostics::format_enumeration;
 use crate::{Db, FxIndexMap, FxOrderMap, Module, ModuleName, Program, declare_lint};
@@ -29,6 +33,7 @@ use std::fmt::Formatter;
 
 /// Registers all known type check lints.
 pub(crate) fn register_lints(registry: &mut LintRegistryBuilder) {
+    registry.register_lint(&AMBIGUOUS_PROTOCOL_MEMBER);
     registry.register_lint(&CALL_NON_CALLABLE);
     registry.register_lint(&POSSIBLY_UNBOUND_IMPLICIT_CALL);
     registry.register_lint(&CONFLICTING_ARGUMENT_FORMS);
@@ -424,7 +429,7 @@ declare_lint! {
 
 declare_lint! {
     /// ## What it does
-    /// Checks for invalidly defined protocol classes.
+    /// Checks for protocol classes that will raise `TypeError` at runtime.
     ///
     /// ## Why is this bad?
     /// An invalidly defined protocol class may lead to the type checker inferring
@@ -450,6 +455,41 @@ declare_lint! {
     }
 }
 
+declare_lint! {
+    /// ## What it does
+    /// Checks for protocol classes with members that will lead to ambiguous interfaces.
+    ///
+    /// ## Why is this bad?
+    /// Assigning to an undeclared variable in a protocol class leads to an ambiguous
+    /// interface which may lead to the type checker inferring unexpected things. It's
+    /// recommended to ensure that all members of a protocol class are explicitly declared.
+    ///
+    /// ## Examples
+    ///
+    /// ```py
+    /// from typing import Protocol
+    ///
+    /// class BaseProto(Protocol):
+    ///     a: int                               # fine (explicitly declared as `int`)
+    ///     def method_member(self) -> int: ...  # fine: a method definition using `def` is considered a declaration
+    ///     c = "some variable"                  # error: no explicit declaration, leading to ambiguity
+    ///     b = method_member                    # error: no explicit declaration, leading to ambiguity
+    ///
+    ///     # error: this creates implicit assignments of `d` and `e` in the protocol class body.
+    ///     # Were they really meant to be considered protocol members?
+    ///     for d, e in enumerate(range(42)):
+    ///         pass
+    ///
+    /// class SubProto(BaseProto, Protocol):
+    ///     a = 42  # fine (declared in superclass)
+    /// ```
+    pub(crate) static AMBIGUOUS_PROTOCOL_MEMBER = {
+        summary: "detects protocol classes with ambiguous interfaces",
+        status: LintStatus::preview("1.0.0"),
+        default_level: Level::Warn,
+    }
+}
+
 declare_lint! {
     /// ## What it does
     /// Checks for invalidly defined `NamedTuple` classes.
@@ -2456,6 +2496,95 @@ pub(crate) fn report_attempted_protocol_instantiation(
     diagnostic.sub(class_def_diagnostic);
 }
 
+pub(crate) fn report_undeclared_protocol_member(
+    context: &InferContext,
+    definition: Definition,
+    protocol_class: ProtocolClassLiteral,
+    class_symbol_table: &PlaceTable,
+) {
+    /// We want to avoid suggesting an annotation for e.g. `x = None`,
+    /// because the user almost certainly doesn't want to write `x: None = None`.
+    /// We also want to avoid suggesting invalid syntax such as `x: <class 'int'> = int`.
+    fn should_give_hint<'db>(db: &'db dyn Db, ty: Type<'db>) -> bool {
+        let class = match ty {
+            Type::ProtocolInstance(ProtocolInstanceType {
+                inner: Protocol::FromClass(_),
+                ..
+            }) => return true,
+            Type::SubclassOf(subclass_of) => match subclass_of.subclass_of() {
+                SubclassOfInner::Class(class) => class,
+                SubclassOfInner::Dynamic(DynamicType::Any) => return true,
+                SubclassOfInner::Dynamic(_) => return false,
+            },
+            Type::NominalInstance(instance) => instance.class(db),
+            _ => return false,
+        };
+
+        !matches!(
+            class.known(db),
+            Some(KnownClass::NoneType | KnownClass::EllipsisType)
+        )
+    }
+
+    let db = context.db();
+
+    let Some(builder) = context.report_lint(
+        &AMBIGUOUS_PROTOCOL_MEMBER,
+        definition.full_range(db, context.module()),
+    ) else {
+        return;
+    };
+
+    let ScopedPlaceId::Symbol(symbol_id) = definition.place(db) else {
+        return;
+    };
+
+    let symbol_name = class_symbol_table.symbol(symbol_id).name();
+    let class_name = protocol_class.name(db);
+
+    let mut diagnostic = builder
+        .into_diagnostic("Cannot assign to undeclared variable in the body of a protocol class");
+
+    if definition.kind(db).is_unannotated_assignment() {
+        let binding_type = binding_type(db, definition);
+
+        let suggestion = binding_type
+            .literal_fallback_instance(db)
+            .unwrap_or(binding_type);
+
+        if should_give_hint(db, suggestion) {
+            diagnostic.set_primary_message(format_args!(
+                "Consider adding an annotation, e.g. `{symbol_name}: {} = ...`",
+                suggestion.display(db)
+            ));
+        } else {
+            diagnostic.set_primary_message(format_args!(
+                "Consider adding an annotation for `{symbol_name}`"
+            ));
+        }
+    } else {
+        diagnostic.set_primary_message(format_args!(
+            "`{symbol_name}` is not declared as a protocol member"
+        ));
+    }
+
+    let mut class_def_diagnostic = SubDiagnostic::new(
+        SubDiagnosticSeverity::Info,
+        "Assigning to an undeclared variable in a protocol class \
+    leads to an ambiguous interface",
+    );
+    class_def_diagnostic.annotate(
+        Annotation::primary(protocol_class.header_span(db))
+            .message(format_args!("`{class_name}` declared as a protocol here",)),
+    );
+    diagnostic.sub(class_def_diagnostic);
+
+    diagnostic.info(format_args!(
+        "No declarations found for `{symbol_name}` \
+        in the body of `{class_name}` or any of its superclasses"
+    ));
+}
+
 pub(crate) fn report_duplicate_bases(
     context: &InferContext,
     class: ClassLiteral,

crates/ty_python_semantic/src/types/infer.rs

@@ -1434,6 +1434,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     }
                 }
             }
+
+            if let Some(protocol) = class.into_protocol_class(self.db()) {
+                protocol.validate_members(&self.context, self.index);
+            }
         }
     }
 

crates/ty_python_semantic/src/types/protocol_class.rs

@@ -7,7 +7,10 @@ use ruff_python_ast::name::Name;
 use rustc_hash::FxHashMap;
 
 use super::TypeVarVariance;
-use crate::semantic_index::place_table;
+use crate::semantic_index::place::ScopedPlaceId;
+use crate::semantic_index::{SemanticIndex, place_table};
+use crate::types::context::InferContext;
+use crate::types::diagnostic::report_undeclared_protocol_member;
 use crate::{
     Db, FxOrderSet,
     place::{Boundness, Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations},
@@ -55,6 +58,59 @@ impl<'db> ProtocolClassLiteral<'db> {
         self.known_function_decorators(db)
             .contains(&KnownFunction::RuntimeCheckable)
     }
+
+    /// Iterate through the body of the protocol class. Check that all definitions
+    /// in the protocol class body are either explicitly declared directly in the
+    /// class body, or are declared in a superclass of the protocol class.
+    pub(super) fn validate_members(self, context: &InferContext, index: &SemanticIndex<'db>) {
+        let db = context.db();
+        let interface = self.interface(db);
+        let class_place_table = index.place_table(self.body_scope(db).file_scope_id(db));
+
+        for (symbol_id, mut bindings_iterator) in
+            use_def_map(db, self.body_scope(db)).all_end_of_scope_symbol_bindings()
+        {
+            let symbol_name = class_place_table.symbol(symbol_id).name();
+
+            if !interface.includes_member(db, symbol_name) {
+                continue;
+            }
+
+            let has_declaration = self
+                .iter_mro(db, None)
+                .filter_map(ClassBase::into_class)
+                .any(|superclass| {
+                    let superclass_scope = superclass.class_literal(db).0.body_scope(db);
+                    let Some(scoped_symbol_id) =
+                        place_table(db, superclass_scope).symbol_id(symbol_name)
+                    else {
+                        return false;
+                    };
+                    !place_from_declarations(
+                        db,
+                        index
+                            .use_def_map(superclass_scope.file_scope_id(db))
+                            .end_of_scope_declarations(ScopedPlaceId::Symbol(scoped_symbol_id)),
+                    )
+                    .into_place_and_conflicting_declarations()
+                    .0
+                    .place
+                    .is_unbound()
+                });
+
+            if has_declaration {
+                continue;
+            }
+
+            let Some(first_definition) =
+                bindings_iterator.find_map(|binding| binding.binding.definition())
+            else {
+                continue;
+            };
+
+            report_undeclared_protocol_member(context, first_definition, self, class_place_table);
+        }
+    }
 }
 
 impl<'db> Deref for ProtocolClassLiteral<'db> {
@@ -147,6 +203,10 @@ impl<'db> ProtocolInterface<'db> {
         })
     }
 
+    pub(super) fn includes_member(self, db: &'db dyn Db, name: &str) -> bool {
+        self.inner(db).contains_key(name)
+    }
+
     pub(super) fn instance_member(self, db: &'db dyn Db, name: &str) -> PlaceAndQualifiers<'db> {
         self.member_by_name(db, name)
             .map(|member| PlaceAndQualifiers {