[ty] Add a subdiagnostic if invalid-return-type is emitted on a method with an empty body on a non-protocol subclass of a protocol class (#18243)

crates/ty_python_semantic/resources/mdtest/function/return_type.md

@@ -397,3 +397,19 @@ async def i() -> typing.AsyncIterable:
 async def j() -> str:  # error: [invalid-return-type]
     yield 42
 ```
+
+## Diagnostics for `invalid-return-type` on non-protocol subclasses of protocol classes
+
+<!-- snapshot-diagnostics -->
+
+We emit a nice subdiagnostic in this situation explaining the probable error here:
+
+```py
+from typing_extensions import Protocol
+
+class Abstract(Protocol):
+    def method(self) -> str: ...
+
+class Concrete(Abstract):
+    def method(self) -> str: ...  # error: [invalid-return-type]
+```

crates/ty_python_semantic/resources/mdtest/snapshots/return_type.md_-_Function_return_type_-_Diagnostics_for_`inv…_(35563a74094b14d5).snap

@@ -0,0 +1,48 @@
+---
+source: crates/ty_test/src/lib.rs
+expression: snapshot
+---
+---
+mdtest name: return_type.md - Function return type - Diagnostics for `invalid-return-type` on non-protocol subclasses of protocol classes
+mdtest path: crates/ty_python_semantic/resources/mdtest/function/return_type.md
+---
+
+# Python source files
+
+## mdtest_snippet.py
+
+```
+1 | from typing_extensions import Protocol
+2 | 
+3 | class Abstract(Protocol):
+4 |     def method(self) -> str: ...
+5 | 
+6 | class Concrete(Abstract):
+7 |     def method(self) -> str: ...  # error: [invalid-return-type]
+```
+
+# Diagnostics
+
+```
+error[invalid-return-type]: Function can implicitly return `None`, which is not assignable to return type `str`
+ --> src/mdtest_snippet.py:7:25
+  |
+6 | class Concrete(Abstract):
+7 |     def method(self) -> str: ...  # error: [invalid-return-type]
+  |                         ^^^
+  |
+info: Only functions in stub files, methods on protocol classes, or methods with `@abstractmethod` are permitted to have empty bodies
+info: Class `Concrete` has `typing.Protocol` in its MRO, but it is not a protocol class
+info: Only classes that directly inherit from `typing.Protocol` or `typing_extensions.Protocol` are considered protocol classes
+ --> src/mdtest_snippet.py:6:7
+  |
+4 |     def method(self) -> str: ...
+5 |
+6 | class Concrete(Abstract):
+  |       ^^^^^^^^^^^^^^^^^^ `Protocol` not present in `Concrete`'s immediate bases
+7 |     def method(self) -> str: ...  # error: [invalid-return-type]
+  |
+info: See https://typing.python.org/en/latest/spec/protocol.html#
+info: rule `invalid-return-type` is enabled by default
+
+```

crates/ty_python_semantic/src/semantic_index/symbol.rs

@@ -546,6 +546,13 @@ impl NodeWithScopeKind {
         }
     }
 
+    pub(crate) const fn as_class(&self) -> Option<&ast::StmtClassDef> {
+        match self {
+            Self::Class(class) => Some(class.node()),
+            _ => None,
+        }
+    }
+
     pub fn expect_function(&self) -> &ast::StmtFunctionDef {
         self.as_function().expect("expected function")
     }

crates/ty_python_semantic/src/types.rs

@@ -1,4 +1,4 @@
-use infer::enclosing_class_symbol;
+use infer::nearest_enclosing_class;
 use itertools::Either;
 
 use std::slice::Iter;
@@ -4906,7 +4906,7 @@ impl<'db> Type<'db> {
 
                 KnownInstanceType::TypingSelf => {
                     let index = semantic_index(db, scope_id.file(db));
-                    let Some(class_ty) = enclosing_class_symbol(db, index, scope_id) else {
+                    let Some(class) = nearest_enclosing_class(db, index, scope_id) else {
                         return Err(InvalidTypeExpressionError {
                             fallback_type: Type::unknown(),
                             invalid_expressions: smallvec::smallvec![
@@ -4914,24 +4914,13 @@ impl<'db> Type<'db> {
                             ],
                         });
                     };
-                    let Some(TypeDefinition::Class(class_def)) = class_ty.definition(db) else {
+                    let instance = Type::ClassLiteral(class)
-                        debug_assert!(
+                        .to_instance(db)
-                            false,
+                        .expect("enclosing_class_symbol must return type that can be instantiated");
-                            "enclosing_class_symbol must return a type with class definition"
-                        );
-                        return Ok(Type::unknown());
-                    };
-                    let Some(instance) = class_ty.to_instance(db) else {
-                        debug_assert!(
-                            false,
-                            "enclosing_class_symbol must return type that can be instantiated"
-                        );
-                        return Ok(Type::unknown());
-                    };
                     Ok(Type::TypeVar(TypeVarInstance::new(
                         db,
                         ast::name::Name::new("Self"),
-                        Some(class_def),
+                        Some(class.definition(db)),
                         Some(TypeVarBoundOrConstraints::UpperBound(instance)),
                         TypeVarVariance::Invariant,
                         None,

crates/ty_python_semantic/src/types/diagnostic.rs

@@ -1,6 +1,6 @@
 use super::context::InferContext;
 use super::mro::DuplicateBaseError;
-use super::{ClassLiteral, KnownClass};
+use super::{ClassBase, ClassLiteral, KnownClass};
 use crate::db::Db;
 use crate::lint::{Level, LintRegistryBuilder, LintStatus};
 use crate::suppression::FileSuppressionId;
@@ -1624,14 +1624,51 @@ pub(super) fn report_implicit_return_type(
     context: &InferContext,
     range: impl Ranged,
     expected_ty: Type,
+    has_empty_body: bool,
+    enclosing_class_of_method: Option<ClassLiteral>,
 ) {
     let Some(builder) = context.report_lint(&INVALID_RETURN_TYPE, range) else {
         return;
     };
-    builder.into_diagnostic(format_args!(
+    let db = context.db();
+    let mut diagnostic = builder.into_diagnostic(format_args!(
         "Function can implicitly return `None`, which is not assignable to return type `{}`",
-        expected_ty.display(context.db())
+        expected_ty.display(db)
     ));
+    if !has_empty_body {
+        return;
+    }
+    let Some(class) = enclosing_class_of_method else {
+        return;
+    };
+    if class
+        .iter_mro(db, None)
+        .any(|base| matches!(base, ClassBase::Protocol(_)))
+    {
+        diagnostic.info(
+            "Only functions in stub files, methods on protocol classes, \
+            or methods with `@abstractmethod` are permitted to have empty bodies",
+        );
+        diagnostic.info(format_args!(
+            "Class `{}` has `typing.Protocol` in its MRO, but it is not a protocol class",
+            class.name(db)
+        ));
+
+        let mut sub_diagnostic = SubDiagnostic::new(
+            Severity::Info,
+            "Only classes that directly inherit from `typing.Protocol` \
+            or `typing_extensions.Protocol` are considered protocol classes",
+        );
+        sub_diagnostic.annotate(
+            Annotation::primary(class.header_span(db)).message(format_args!(
+                "`Protocol` not present in `{class}`'s immediate bases",
+                class = class.name(db)
+            )),
+        );
+        diagnostic.sub(sub_diagnostic);
+
+        diagnostic.info("See https://typing.python.org/en/latest/spec/protocol.html#");
+    }
 }
 
 pub(super) fn report_invalid_type_checking_constant(context: &InferContext, node: AnyNodeRef) {

crates/ty_python_semantic/src/types/infer.rs

@@ -333,25 +333,26 @@ fn unpack_cycle_initial<'db>(_db: &'db dyn Db, _unpack: Unpack<'db>) -> UnpackRe
 /// Returns the type of the nearest enclosing class for the given scope.
 ///
 /// This function walks up the ancestor scopes starting from the given scope,
-/// and finds the closest class definition.
+/// and finds the closest class definition. This is different to the behaviour of
+/// [`TypeInferenceBuilder::class_context_of_current_method`], which will only return
+/// `Some(class)` if either the immediate parent scope is a class OR the immediate parent
+/// scope is a type-parameters scope and the grandparent scope is a class.
 ///
-/// Returns `None` if no enclosing class is found.a
+/// Returns `None` if no enclosing class is found.
-pub(crate) fn enclosing_class_symbol<'db>(
+pub(crate) fn nearest_enclosing_class<'db>(
     db: &'db dyn Db,
     semantic: &SemanticIndex<'db>,
     scope: ScopeId,
-) -> Option<Type<'db>> {
+) -> Option<ClassLiteral<'db>> {
     semantic
         .ancestor_scopes(scope.file_scope_id(db))
         .find_map(|(_, ancestor_scope)| {
-            if let NodeWithScopeKind::Class(class) = ancestor_scope.node() {
+            let class = ancestor_scope.node().as_class()?;
-                let definition = semantic.expect_single_definition(class.node());
+            let definition = semantic.expect_single_definition(class);
-                let result = infer_definition_types(db, definition);
+            infer_definition_types(db, definition)
-
+                .declaration_type(definition)
-                Some(result.declaration_type(definition).inner_type())
+                .inner_type()
-            } else {
+                .into_class_literal()
-                None
-            }
         })
 }
 
@@ -1691,43 +1692,39 @@ impl<'db> TypeInferenceBuilder<'db> {
         self.infer_annotation_expression(&type_alias.value, DeferredExpressionState::Deferred);
     }
 
-    /// Returns `true` if the current scope is the function body scope of a method of a protocol
+    /// If the current scope is a method inside an enclosing class,
-    /// (that is, a class which directly inherits `typing.Protocol`.)
+    /// return `Some(class)` where `class` represents the enclosing class.
-    fn in_protocol_class(&self) -> bool {
+    ///
+    /// If the current scope is not a method inside an enclosing class,
+    /// return `None`.
+    ///
+    /// Note that this method will only return `Some` if the immediate parent scope
+    /// is a class scope OR the immediate parent scope is an annotation scope
+    /// and the grandparent scope is a class scope. This means it has different
+    /// behaviour to the [`nearest_enclosing_class`] function.
+    fn class_context_of_current_method(&self) -> Option<ClassLiteral<'db>> {
         let current_scope_id = self.scope().file_scope_id(self.db());
         let current_scope = self.index.scope(current_scope_id);
-        let Some(parent_scope_id) = current_scope.parent() else {
+        let parent_scope_id = current_scope.parent()?;
-            return false;
-        };
         let parent_scope = self.index.scope(parent_scope_id);
 
         let class_scope = match parent_scope.kind() {
             ScopeKind::Class => parent_scope,
             ScopeKind::Annotation => {
-                let Some(class_scope_id) = parent_scope.parent() else {
+                let class_scope_id = parent_scope.parent()?;
-                    return false;
-                };
                 let potentially_class_scope = self.index.scope(class_scope_id);
 
                 match potentially_class_scope.kind() {
                     ScopeKind::Class => potentially_class_scope,
-                    _ => return false,
+                    _ => return None,
                 }
             }
-            _ => return false,
+            _ => return None,
         };
 
-        let NodeWithScopeKind::Class(node_ref) = class_scope.node() else {
+        let class_stmt = class_scope.node().as_class()?;
-            return false;
+        let class_definition = self.index.expect_single_definition(class_stmt);
-        };
+        binding_type(self.db(), class_definition).into_class_literal()
-
-        let class_definition = self.index.expect_single_definition(node_ref.node());
-
-        let Type::ClassLiteral(class) = binding_type(self.db(), class_definition) else {
-            return false;
-        };
-
-        class.is_protocol(self.db())
     }
 
     /// Returns `true` if the current scope is the function body scope of a function overload (that
@@ -1789,14 +1786,25 @@ impl<'db> TypeInferenceBuilder<'db> {
                 }
             }
 
-            if (self.in_stub()
+            let has_empty_body =
-                || self.in_function_overload_or_abstractmethod()
+                self.return_types_and_ranges.is_empty() && is_stub_suite(&function.body);
-                || self.in_protocol_class())
+
-                && self.return_types_and_ranges.is_empty()
+            let mut enclosing_class_context = None;
-                && is_stub_suite(&function.body)
+
-            {
+            if has_empty_body {
+                if self.in_stub() {
                     return;
                 }
+                if self.in_function_overload_or_abstractmethod() {
+                    return;
+                }
+                if let Some(class) = self.class_context_of_current_method() {
+                    enclosing_class_context = Some(class);
+                    if class.is_protocol(self.db()) {
+                        return;
+                    }
+                }
+            }
 
             let declared_ty = self.file_expression_type(returns);
 
@@ -1856,7 +1864,13 @@ impl<'db> TypeInferenceBuilder<'db> {
             if use_def.can_implicit_return(self.db())
                 && !Type::none(self.db()).is_assignable_to(self.db(), declared_ty)
             {
-                report_implicit_return_type(&self.context, returns.range(), declared_ty);
+                report_implicit_return_type(
+                    &self.context,
+                    returns.range(),
+                    declared_ty,
+                    has_empty_body,
+                    enclosing_class_context,
+                );
             }
         }
     }
@@ -2139,7 +2153,9 @@ impl<'db> TypeInferenceBuilder<'db> {
                     }
                 } else if (self.in_stub()
                     || self.in_function_overload_or_abstractmethod()
-                    || self.in_protocol_class())
+                    || self
+                        .class_context_of_current_method()
+                        .is_some_and(|class| class.is_protocol(self.db())))
                     && default
                         .as_ref()
                         .is_some_and(|d| d.is_ellipsis_literal_expr())
@@ -5144,7 +5160,7 @@ impl<'db> TypeInferenceBuilder<'db> {
                                             [] => {
                                                 let scope = self.scope();
 
-                                                let Some(enclosing_class) = enclosing_class_symbol(
+                                                let Some(enclosing_class) = nearest_enclosing_class(
                                                     self.db(),
                                                     self.index,
                                                     scope,
@@ -5172,7 +5188,7 @@ impl<'db> TypeInferenceBuilder<'db> {
 
                                                 let bound_super = BoundSuperType::build(
                                                     self.db(),
-                                                    enclosing_class,
+                                                    Type::ClassLiteral(enclosing_class),
                                                     first_param,
                                                 )
                                                 .unwrap_or_else(|err| {