[red-knot] Add AlwaysTruthy and AlwaysFalsy to knot_extensions

crates/red_knot_python_semantic/resources/mdtest/type_api.md

@@ -94,6 +94,27 @@ reveal_type(C.__mro__)
 u: Unknown[str]
 ```
 
+### `AlwaysTruthy` and `AlwaysFalsy`
+
+`AlwaysTruthy` and `AlwaysFalsy` represent the sets of all types whose truthiness is always truthy
+or falsy, respectively.
+
+```py
+from typing_extensions import Literal
+
+from knot_extensions import AlwaysFalsy, AlwaysTruthy, is_subtype_of, static_assert
+
+static_assert(is_subtype_of(Literal[True], AlwaysTruthy))
+static_assert(is_subtype_of(Literal[False], AlwaysFalsy))
+
+static_assert(not is_subtype_of(int, AlwaysFalsy))
+static_assert(not is_subtype_of(str, AlwaysFalsy))
+
+def _(t: AlwaysTruthy, f: AlwaysFalsy):
+    reveal_type(t)  # revealed: AlwaysTruthy
+    reveal_type(f)  # revealed: AlwaysFalsy
+```
+
 ## Static assertions
 
 ### Basics
@@ -350,3 +371,37 @@ def type_of_annotation() -> None:
 # error: "Special form `knot_extensions.TypeOf` expected exactly one type parameter"
 t: TypeOf[int, str, bytes]
 ```
+
+## `LiteralExt`
+
+`LiteralExt` creates (unions of) class, function or slice types.
+
+It cannot be nested and does not mix with `Literal`.
+
+```py
+from typing_extensions import Literal
+
+from knot_extensions import LiteralExt, static_assert, is_subtype_of
+
+class C: ...
+class D: ...
+
+def f(): ...
+def g(): ...
+def _(a: LiteralExt[C], b: LiteralExt[f], c: LiteralExt[1:2:3], flag: bool):
+    reveal_type(a)  # revealed: Literal[C]
+    reveal_type(b)  # revealed: Literal[f]
+    reveal_type(c)  # revealed: slice[Literal[1], Literal[2], Literal[3]]
+
+def _(a: LiteralExt[C, D], b: LiteralExt[C, f, g], c: LiteralExt[D, g, ::4]):
+    reveal_type(a)  # revealed: Literal[C, D]
+    reveal_type(b)  # revealed: Literal[C] | Literal[f, g]
+    reveal_type(c)  # revealed: Literal[D] | Literal[g] | slice[None, None, Literal[4]]
+
+# error: [invalid-type-form]
+d: LiteralExt[LiteralExt[f]]
+# error: [invalid-type-form]
+d: LiteralExt[4]
+# error: [invalid-type-form]
+d: LiteralExt[Literal["5"]]
+```

crates/red_knot_python_semantic/src/types.rs

@@ -2278,6 +2278,8 @@ impl<'db> Type<'db> {
                 fallback_type: Type::unknown(),
             }),
             Type::KnownInstance(KnownInstanceType::Unknown) => Ok(Type::unknown()),
+            Type::KnownInstance(KnownInstanceType::AlwaysTruthy) => Ok(Type::AlwaysTruthy),
+            Type::KnownInstance(KnownInstanceType::AlwaysFalsy) => Ok(Type::AlwaysFalsy),
             _ => Ok(todo_type!(
                 "Unsupported or invalid type in a type expression"
             )),
@@ -2813,12 +2815,18 @@ pub enum KnownInstanceType<'db> {
     TypeAliasType(TypeAliasType<'db>),
     /// The symbol `knot_extensions.Unknown`
     Unknown,
+    /// The symbol `knot_extensions.AlwaysTruthy`
+    AlwaysTruthy,
+    /// The symbol `knot_extensions.AlwaysFalsy`
+    AlwaysFalsy,
     /// The symbol `knot_extensions.Not`
     Not,
     /// The symbol `knot_extensions.Intersection`
     Intersection,
     /// The symbol `knot_extensions.TypeOf`
     TypeOf,
+    /// The symbol `knot_extensions.LiteralExt`
+    LiteralExt,
 
     // Various special forms, special aliases and type qualifiers that we don't yet understand
     // (all currently inferred as TODO in most contexts):
@@ -2874,9 +2882,12 @@ impl<'db> KnownInstanceType<'db> {
             Self::OrderedDict => "OrderedDict",
             Self::ReadOnly => "ReadOnly",
             Self::Unknown => "Unknown",
+            Self::AlwaysTruthy => "AlwaysTruthy",
+            Self::AlwaysFalsy => "AlwaysFalsy",
             Self::Not => "Not",
             Self::Intersection => "Intersection",
             Self::TypeOf => "TypeOf",
+            Self::LiteralExt => "LiteralExt",
         }
     }
 
@@ -2917,9 +2928,12 @@ impl<'db> KnownInstanceType<'db> {
             | Self::ReadOnly
             | Self::TypeAliasType(_)
             | Self::Unknown
+            | Self::AlwaysTruthy
+            | Self::AlwaysFalsy
             | Self::Not
             | Self::Intersection
-            | Self::TypeOf => Truthiness::AlwaysTrue,
+            | Self::TypeOf
+            | Self::LiteralExt => Truthiness::AlwaysTrue,
         }
     }
 
@@ -2960,9 +2974,12 @@ impl<'db> KnownInstanceType<'db> {
             Self::TypeVar(typevar) => typevar.name(db),
             Self::TypeAliasType(_) => "typing.TypeAliasType",
             Self::Unknown => "knot_extensions.Unknown",
+            Self::AlwaysTruthy => "knot_extensions.AlwaysTruthy",
+            Self::AlwaysFalsy => "knot_extensions.AlwaysFalsy",
             Self::Not => "knot_extensions.Not",
             Self::Intersection => "knot_extensions.Intersection",
             Self::TypeOf => "knot_extensions.TypeOf",
+            Self::LiteralExt => "knot_extensions.LiteralExt",
         }
     }
 
@@ -3005,7 +3022,10 @@ impl<'db> KnownInstanceType<'db> {
             Self::TypeOf => KnownClass::SpecialForm,
             Self::Not => KnownClass::SpecialForm,
             Self::Intersection => KnownClass::SpecialForm,
+            Self::LiteralExt => KnownClass::SpecialForm,
             Self::Unknown => KnownClass::Object,
+            Self::AlwaysTruthy => KnownClass::Object,
+            Self::AlwaysFalsy => KnownClass::Object,
         }
     }
 
@@ -3057,9 +3077,12 @@ impl<'db> KnownInstanceType<'db> {
             "NotRequired" => Self::NotRequired,
             "LiteralString" => Self::LiteralString,
             "Unknown" => Self::Unknown,
+            "AlwaysTruthy" => Self::AlwaysTruthy,
+            "AlwaysFalsy" => Self::AlwaysFalsy,
             "Not" => Self::Not,
             "Intersection" => Self::Intersection,
             "TypeOf" => Self::TypeOf,
+            "LiteralExt" => Self::LiteralExt,
             _ => return None,
         };
 
@@ -3109,9 +3132,13 @@ impl<'db> KnownInstanceType<'db> {
             | Self::TypeVar(_) => {
                 matches!(module, KnownModule::Typing | KnownModule::TypingExtensions)
             }
-            Self::Unknown | Self::Not | Self::Intersection | Self::TypeOf => {
-                module.is_knot_extensions()
-            }
+            Self::Unknown
+            | Self::AlwaysTruthy
+            | Self::AlwaysFalsy
+            | Self::Not
+            | Self::Intersection
+            | Self::TypeOf
+            | Self::LiteralExt => module.is_knot_extensions(),
         }
     }
 

crates/red_knot_python_semantic/src/types/class_base.rs

@@ -105,7 +105,10 @@ impl<'db> ClassBase<'db> {
                 | KnownInstanceType::Optional
                 | KnownInstanceType::Not
                 | KnownInstanceType::Intersection
-                | KnownInstanceType::TypeOf => None,
+                | KnownInstanceType::TypeOf
+                | KnownInstanceType::LiteralExt
+                | KnownInstanceType::AlwaysTruthy
+                | KnownInstanceType::AlwaysFalsy => None,
                 KnownInstanceType::Unknown => Some(Self::unknown()),
                 KnownInstanceType::Any => Some(Self::any()),
                 // TODO: Classes inheriting from `typing.Type` et al. also have `Generic` in their MRO

crates/red_knot_python_semantic/src/types/infer.rs

@@ -4539,7 +4539,7 @@ impl<'db> TypeInferenceBuilder<'db> {
 
                         return dunder_getitem_method
                             .call(self.db(), &CallArguments::positional([value_ty, slice_ty]))
-                            .return_ty_result( &self.context, value_node.into())
+                            .return_ty_result(&self.context, value_node.into())
                             .unwrap_or_else(|err| {
                                 self.context.report_lint(
                                     &CALL_NON_CALLABLE,
@@ -5298,6 +5298,39 @@ impl<'db> TypeInferenceBuilder<'db> {
                     argument_type
                 }
             },
+            KnownInstanceType::LiteralExt => {
+                let elements = match arguments_slice {
+                    ast::Expr::Tuple(tuple) => Either::Left(tuple.iter()),
+                    element => Either::Right(std::iter::once(element)),
+                };
+
+                elements
+                    .fold(UnionBuilder::new(self.db()), |builder, element| {
+                        let argument_type = self.infer_expression(element);
+
+                        match argument_type {
+                            Type::ClassLiteral(_)
+                            | Type::FunctionLiteral(_)
+                            | Type::SliceLiteral(_) => {
+                                builder.add(argument_type)
+                            }
+
+                            _ => {
+                                self.context.report_lint(
+                                    &INVALID_TYPE_FORM,
+                                    element.into(),
+                                    format_args!(
+                                        "Special form `{}` expected class, function or slice literals as arguments, got `{}`",
+                                        known_instance.repr(self.db()),
+                                        argument_type.display(self.db()),
+                                    ),
+                                );
+                                builder.add(Type::unknown())
+                            }
+                        }
+                    })
+                    .build()
+            }
 
             // TODO: Generics
             KnownInstanceType::ChainMap => {
@@ -5397,6 +5430,22 @@ impl<'db> TypeInferenceBuilder<'db> {
                 );
                 Type::unknown()
             }
+            KnownInstanceType::AlwaysTruthy | KnownInstanceType::AlwaysFalsy => {
+                self.context.report_lint(
+                    &INVALID_TYPE_FORM,
+                    subscript.into(),
+                    format_args!(
+                        "Special form `{}` expected no type parameter",
+                        known_instance.repr(self.db())
+                    ),
+                );
+
+                if matches!(known_instance, KnownInstanceType::AlwaysTruthy) {
+                    Type::AlwaysTruthy
+                } else {
+                    Type::AlwaysFalsy
+                }
+            }
             KnownInstanceType::LiteralString => {
                 self.context.report_lint(
                     &INVALID_TYPE_FORM,

crates/red_knot_vendored/knot_extensions/knot_extensions.pyi

@@ -5,11 +5,14 @@ def static_assert(condition: object, msg: LiteralString | None = None) -> None:
 
 # Types
 Unknown = object()
+AlwaysTruthy = object()
+AlwaysFalsy = object()
 
 # Special forms
 Not: _SpecialForm
 Intersection: _SpecialForm
 TypeOf: _SpecialForm
+LiteralExt: _SpecialForm
 
 # Predicates on types
 #