Disallow @disjoint_base on TypedDicts and Protocols (#24671)

Author: JelleZijlstra

crates/ty_python_semantic/resources/mdtest/instance_layout_conflict.md

@@ -183,7 +183,7 @@ decorator introduced by this PEP provides a generalised way for type checkers to
 classes.
 
 ```py
-from typing_extensions import disjoint_base
+from typing_extensions import Protocol, TypedDict, disjoint_base
 
 # fmt: off
 
@@ -213,12 +213,16 @@ class G: ...
 
 @disjoint_base
 class H: ...
-
+@disjoint_base  # error: [invalid-typed-dict-header] "`@disjoint_base` cannot be used with `TypedDict` class `Movie`"
+class Movie(TypedDict):
+    name: str
+@disjoint_base  # error: [invalid-protocol] "`@disjoint_base` cannot be used with protocol class `SupportsClose`"
+class SupportsClose(Protocol):
+    def close(self) -> None: ...
 class I(  # error: [instance-layout-conflict]
     G,
     H
 ): ...
-
 # fmt: on
 ```
 

crates/ty_python_semantic/resources/mdtest/snapshots/instance_layout_conf…_-_Tests_for_ty's_`inst…_-_Builtins_with_implic…_(4c3d127986a58f11).snap

@@ -13,7 +13,7 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/instance_layout_conflict
 ## mdtest_snippet.py
 
 ```
- 1 | from typing_extensions import disjoint_base
+ 1 | from typing_extensions import Protocol, TypedDict, disjoint_base
  2 | 
  3 | # fmt: off
  4 | 
@@ -43,15 +43,19 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/instance_layout_conflict
 28 | 
 29 | @disjoint_base
 30 | class H: ...
-31 | 
-32 | class I(  # error: [instance-layout-conflict]
-33 |     G,
-34 |     H
-35 | ): ...
-36 | 
-37 | # fmt: on
-38 | # error: [invalid-generic-class]
-39 | class Foo(range, str): ...  # error: [subclass-of-final-class]
+31 | @disjoint_base  # error: [invalid-typed-dict-header] "`@disjoint_base` cannot be used with `TypedDict` class `Movie`"
+32 | class Movie(TypedDict):
+33 |     name: str
+34 | @disjoint_base  # error: [invalid-protocol] "`@disjoint_base` cannot be used with protocol class `SupportsClose`"
+35 | class SupportsClose(Protocol):
+36 |     def close(self) -> None: ...
+37 | class I(  # error: [instance-layout-conflict]
+38 |     G,
+39 |     H
+40 | ): ...
+41 | # fmt: on
+42 | # error: [invalid-generic-class]
+43 | class Foo(range, str): ...  # error: [subclass-of-final-class]
 ```
 
 # Diagnostics
@@ -144,33 +148,53 @@ info: Two classes cannot coexist in a class's MRO if their instances have incomp
 
 ```
 
+```
+error[invalid-typed-dict-header]: `@disjoint_base` cannot be used with `TypedDict` class `Movie`
+  --> src/mdtest_snippet.py:31:1
+   |
+31 | @disjoint_base  # error: [invalid-typed-dict-header] "`@disjoint_base` cannot be used with `TypedDict` class `Movie`"
+   | ^^^^^^^^^^^^^^
+   |
+
+```
+
+```
+error[invalid-protocol]: `@disjoint_base` cannot be used with protocol class `SupportsClose`
+  --> src/mdtest_snippet.py:34:1
+   |
+34 | @disjoint_base  # error: [invalid-protocol] "`@disjoint_base` cannot be used with protocol class `SupportsClose`"
+   | ^^^^^^^^^^^^^^
+   |
+
+```
+
 ```
 error[instance-layout-conflict]: Class will raise `TypeError` at runtime due to incompatible bases
-  --> src/mdtest_snippet.py:32:7
+  --> src/mdtest_snippet.py:37:7
    |
-32 |   class I(  # error: [instance-layout-conflict]
+37 |   class I(  # error: [instance-layout-conflict]
    |  _______^
-33 | |     G,
-34 | |     H
-35 | | ): ...
+38 | |     G,
+39 | |     H
+40 | | ): ...
    | |_^ Bases `G` and `H` cannot be combined in multiple inheritance
    |
 info: Two classes cannot coexist in a class's MRO if their instances have incompatible memory layouts
-  --> src/mdtest_snippet.py:33:5
+  --> src/mdtest_snippet.py:38:5
    |
-33 |     G,
+38 |     G,
    |     - `G` instances have a distinct memory layout because of the way `G` is implemented in a C extension
-34 |     H
+39 |     H
    |     - `H` instances have a distinct memory layout because of the way `H` is implemented in a C extension
    |
 
 ```
 
 ```
 error[invalid-generic-class]: Inconsistent type arguments for `Sequence` among class bases
-  --> src/mdtest_snippet.py:39:7
+  --> src/mdtest_snippet.py:43:7
    |
-39 | class Foo(range, str): ...  # error: [subclass-of-final-class]
+43 | class Foo(range, str): ...  # error: [subclass-of-final-class]
    |       ^^^^-----^^---^
    |           |      |
    |           |      Later class base inherits from `Sequence[str]`
@@ -181,9 +205,9 @@ error[invalid-generic-class]: Inconsistent type arguments for `Sequence` among c
 
 ```
 error[subclass-of-final-class]: Class `Foo` cannot inherit from final class `range`
-  --> src/mdtest_snippet.py:39:11
+  --> src/mdtest_snippet.py:43:11
    |
-39 | class Foo(range, str): ...  # error: [subclass-of-final-class]
+43 | class Foo(range, str): ...  # error: [subclass-of-final-class]
    |           ^^^^^
    |
 

crates/ty_python_semantic/src/types/class/static_literal.rs

@@ -490,6 +490,8 @@ impl<'db> StaticClassLiteral<'db> {
         if self
             .known_function_decorators(db)
             .contains(&KnownFunction::DisjointBase)
+            && !self.is_typed_dict(db)
+            && !self.is_protocol(db)
         {
             Some(DisjointBase::due_to_decorator(self))
         } else if SlotsKind::from(db, self) == SlotsKind::NotEmpty {

crates/ty_python_semantic/src/types/infer/builder/post_inference/static_class.rs

@@ -53,10 +53,10 @@ use crate::{
 use ty_python_core::{SemanticIndex, definition::DefinitionKind, scope::ScopeId};
 
 /// Iterate over all static class definitions (created using `class` statements) to check that
-/// the definition will not cause an exception to be raised at runtime. This needs to be done
-/// after most other types in the scope have been inferred, due to the fact that base classes
-/// can be deferred. If it looks like a class definition is invalid in some way, issue a
-/// diagnostic.
+/// the definition is semantically valid and will not cause an exception to be raised at runtime.
+/// This needs to be done after most other types in the scope have been inferred, due to the fact
+/// that base classes can be deferred. If it looks like a class definition is invalid in some way,
+/// issue a diagnostic.
 ///
 /// Note: Dynamic classes created via `type()` calls are checked separately during type
 /// inference of the call expression.
@@ -142,6 +142,30 @@ pub(crate) fn check_static_class_definitions<'db>(
 
     let is_protocol = class.is_protocol(db);
 
+    if let Some(disjoint_base_decorator) = class_node.decorator_list.iter().find(|decorator| {
+        file_expression_type(&decorator.expression)
+            .as_function_literal()
+            .is_some_and(|function| function.is_known(db, KnownFunction::DisjointBase))
+    }) {
+        if class_kind == Some(CodeGeneratorKind::TypedDict) {
+            if let Some(builder) =
+                context.report_lint(&INVALID_TYPED_DICT_HEADER, disjoint_base_decorator)
+            {
+                builder.into_diagnostic(format_args!(
+                    "`@disjoint_base` cannot be used with `TypedDict` class `{}`",
+                    class.name(db),
+                ));
+            }
+        } else if is_protocol
+            && let Some(builder) = context.report_lint(&INVALID_PROTOCOL, disjoint_base_decorator)
+        {
+            builder.into_diagnostic(format_args!(
+                "`@disjoint_base` cannot be used with protocol class `{}`",
+                class.name(db),
+            ));
+        }
+    }
+
     // Check for invalid `@dataclass` applications.
     if class.dataclass_params(db).is_some() {
         if class.has_named_tuple_class_in_mro(db) {