[red-knot] Handle special case returning NotImplemented (#17034)

## Summary

Closes #16661

This PR includes two changes:

- `NotImplementedType` is now a member of `KnownClass`
- We skip `is_assignable_to` checks for `NotImplemented` when checking
return types

### Limitation

```py
def f(cond: bool) -> int:
    return 1 if cond else NotImplemented
```

The implementation covers cases where `NotImplemented` appears inside a
`Union`.
However, for more complex types (ex. `Intersection`) it will not worked.
In my opinion, supporting such complexity is unnecessary at this point.

## Test Plan

Two `mdtest` files were updated:

- `mdtest/function/return_type.md`
- `mdtest/type_properties/is_singleton.md`

To test `KnownClass`, run:
```bash
cargo test -p red_knot_python_semantic -- types::class::
```

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

@@ -269,3 +269,45 @@ def f(cond: bool) -> int:
     if cond:
         return 2
 ```
+
+## NotImplemented
+
+`NotImplemented` is a special symbol in Python. It is commonly used to control the fallback behavior
+of special dunder methods. You can find more details in the
+[documentation](https://docs.python.org/3/library/numbers.html#implementing-the-arithmetic-operations).
+
+```py
+from __future__ import annotations
+
+class A:
+    def __add__(self, o: A) -> A:
+        return NotImplemented
+```
+
+However, as shown below, `NotImplemented` should not cause issues with the declared return type.
+
+```py
+def f() -> int:
+    return NotImplemented
+
+def f(cond: bool) -> int:
+    if cond:
+        return 1
+    else:
+        return NotImplemented
+
+def f(x: int) -> int | str:
+    if x < 0:
+        return -1
+    elif x == 0:
+        return NotImplemented
+    else:
+        return "test"
+
+def f(cond: bool) -> str:
+    return "hello" if cond else NotImplemented
+
+def f(cond: bool) -> int:
+    # error: [invalid-return-type] "Object of type `Literal["hello"]` is not assignable to return type `int`"
+    return "hello" if cond else NotImplemented
+```

crates/red_knot_python_semantic/resources/mdtest/type_properties/is_singleton.md

@@ -72,7 +72,6 @@ python-version = "3.9"
 ```
 
 ```py
-import sys
 from knot_extensions import is_singleton, static_assert
 
 static_assert(is_singleton(Ellipsis.__class__))
@@ -95,3 +94,42 @@ from knot_extensions import static_assert, is_singleton
 
 static_assert(is_singleton(types.EllipsisType))
 ```
+
+## `builtins.NotImplemented` / `types.NotImplementedType`
+
+### All Python versions
+
+Just like `Ellipsis`, the type of `NotImplemented` was not exposed on Python \<3.10. However, we
+still recognize the type as a singleton in all Python versions.
+
+```toml
+[environment]
+python-version = "3.9"
+```
+
+```py
+from knot_extensions import is_singleton, static_assert
+
+static_assert(is_singleton(NotImplemented.__class__))
+```
+
+### Python 3.10+
+
+On Python 3.10+, the standard library exposes the type of `NotImplemented` as
+`types.NotImplementedType`. We also recognize this as a singleton type when it is referenced
+directly:
+
+```toml
+[environment]
+python-version = "3.10"
+```
+
+```py
+import types
+from knot_extensions import static_assert, is_singleton
+
+# TODO: types.NotImplementedType is a TypeAlias of builtins._NotImplementedType
+# Once TypeAlias support is added, it should satisfy `is_singleton`
+reveal_type(types.NotImplementedType)  # revealed: Unknown | Literal[_NotImplementedType]
+static_assert(not is_singleton(types.NotImplementedType))
+```

crates/red_knot_python_semantic/src/types.rs

@@ -315,6 +315,14 @@ impl<'db> Type<'db> {
             .is_some_and(|instance| instance.class().is_known(db, KnownClass::NoneType))
     }
 
+    pub fn is_notimplemented(&self, db: &'db dyn Db) -> bool {
+        self.into_instance().is_some_and(|instance| {
+            instance
+                .class()
+                .is_known(db, KnownClass::NotImplementedType)
+        })
+    }
+
     pub fn is_object(&self, db: &'db dyn Db) -> bool {
         self.into_instance()
             .is_some_and(|instance| instance.class().is_object(db))
@@ -4975,6 +4983,10 @@ impl<'db> UnionType<'db> {
         Self::from_elements(db, self.elements(db).iter().map(transform_fn))
     }
 
+    pub fn filter(&self, db: &'db dyn Db, filter_fn: impl FnMut(&&Type<'db>) -> bool) -> Type<'db> {
+        Self::from_elements(db, self.elements(db).iter().filter(filter_fn))
+    }
+
     pub(crate) fn map_with_boundness(
         self,
         db: &'db dyn Db,

crates/red_knot_python_semantic/src/types/class.rs

@@ -862,6 +862,7 @@ pub enum KnownClass {
     // Exposed as `types.EllipsisType` on Python >=3.10;
     // backported as `builtins.ellipsis` by typeshed on Python <=3.9
     EllipsisType,
+    NotImplementedType,
 }
 
 impl<'db> KnownClass {
@@ -929,6 +930,10 @@ impl<'db> KnownClass {
             | Self::Sized
             | Self::Enum
             | Self::Super
+            // Evaluating `NotImplementedType` in a boolean context was deprecated in Python 3.9
+            // and raises a `TypeError` in Python >=3.14
+            // (see https://docs.python.org/3/library/constants.html#NotImplemented)
+            | Self::NotImplementedType
             | Self::Classmethod => Truthiness::Ambiguous,
         }
     }
@@ -994,6 +999,7 @@ impl<'db> KnownClass {
                     "ellipsis"
                 }
             }
+            Self::NotImplementedType => "_NotImplementedType",
         }
     }
 
@@ -1167,6 +1173,7 @@ impl<'db> KnownClass {
                     KnownModule::Builtins
                 }
             }
+            Self::NotImplementedType => KnownModule::Builtins,
             Self::ChainMap
             | Self::Counter
             | Self::DefaultDict
@@ -1182,7 +1189,8 @@ impl<'db> KnownClass {
             | Self::NoDefaultType
             | Self::VersionInfo
             | Self::EllipsisType
-            | Self::TypeAliasType => true,
+            | Self::TypeAliasType
+            | Self::NotImplementedType => true,
 
             Self::Bool
             | Self::Object
@@ -1231,13 +1239,13 @@ impl<'db> KnownClass {
     ///
     /// A singleton class is a class where it is known that only one instance can ever exist at runtime.
     pub(super) const fn is_singleton(self) -> bool {
-        // TODO there are other singleton types (NotImplementedType -- any others?)
         match self {
             Self::NoneType
             | Self::EllipsisType
             | Self::NoDefaultType
             | Self::VersionInfo
-            | Self::TypeAliasType => true,
+            | Self::TypeAliasType
+            | Self::NotImplementedType => true,
 
             Self::Bool
             | Self::Object
@@ -1340,6 +1348,9 @@ impl<'db> KnownClass {
             "EllipsisType" if Program::get(db).python_version(db) >= PythonVersion::PY310 => {
                 Self::EllipsisType
             }
+            "_NotImplementedType" if Program::get(db).python_version(db) <= PythonVersion::PY39 => {
+                Self::NotImplementedType
+            }
             _ => return None,
         };
 
@@ -1385,6 +1396,7 @@ impl<'db> KnownClass {
             | Self::MethodWrapperType
             | Self::Enum
             | Self::Super
+            | Self::NotImplementedType
             | Self::WrapperDescriptorType => module == self.canonical_module(db),
             Self::NoneType => matches!(module, KnownModule::Typeshed | KnownModule::Types),
             Self::SpecialForm

crates/red_knot_python_semantic/src/types/infer.rs

@@ -1265,9 +1265,21 @@ impl<'db> TypeInferenceBuilder<'db> {
             {
                 return;
             }
+
             for invalid in self
                 .return_types_and_ranges
                 .iter()
+                .copied()
+                .filter_map(|ty_range| match ty_range.ty {
+                    // We skip `is_assignable_to` checks for `NotImplemented`,
+                    // so we remove it beforehand.
+                    Type::Union(union) => Some(TypeAndRange {
+                        ty: union.filter(self.db(), |ty| !ty.is_notimplemented(self.db())),
+                        range: ty_range.range,
+                    }),
+                    ty if ty.is_notimplemented(self.db()) => None,
+                    _ => Some(ty_range),
+                })
                 .filter(|ty_range| !ty_range.ty.is_assignable_to(self.db(), declared_ty))
             {
                 report_invalid_return_type(