[red-knot] Support assert_type

crates/red_knot_python_semantic/resources/mdtest/directives/assert_type.md

@@ -0,0 +1,123 @@
+from typing import reveal_type
+
+# `assert_type`
+
+## Basic
+
+```py
+from typing_extensions import assert_type
+
+def _(x: int):
+    assert_type(x, int)  # fine
+    assert_type(x, str)  # error: [type-assertion-failure]
+```
+
+## Narrowing
+
+The asserted type is checked against the inferred type, not the declared type.
+
+```py
+from typing_extensions import assert_type
+
+def _(x: int | str):
+    if isinstance(x, int):
+        reveal_type(x)  # revealed: int
+        assert_type(x, int)  # fine
+```
+
+## Equivalence
+
+The actual type must match the asserted type precisely.
+
+```py
+from typing import Any, Type, Union
+from typing_extensions import assert_type
+
+# Subtype does not count
+def _(x: bool):
+    assert_type(x, int)  # error: [type-assertion-failure]
+
+def _(a: type[int], b: type[Any]):
+    assert_type(a, type[Any])  # error: [type-assertion-failure]
+    assert_type(b, type[int])  # error: [type-assertion-failure]
+
+# The expression constructing the type is not taken into account
+def _(a: type[int]):
+    # TODO: Infer the second argument as a type expression
+    assert_type(a, Type[int])  # error: [type-assertion-failure]
+```
+
+## Gradual types
+
+```py
+from typing_extensions import Literal, assert_type
+
+# Any and Unknown are considered equivalent
+def _(a):
+    reveal_type(a)  # revealed: Unknown
+    assert_type(a, Any)  # fine
+
+def _(b: type[Literal]):  # TODO: Should be invalid
+    # TODO: Should be `type[Unknown]`
+    reveal_type(b)  # revealed: @Todo(unsupported type[X] special form)
+
+    # TODO: Infer the second argument as a type expression
+    # Should be fine
+    assert_type(b, type[Any])  # error: [type-assertion-failure]
+```
+
+## Tuples
+
+Tuple types with the same elements are the same.
+
+```py
+from typing_extensions import assert_type
+
+def _(a: tuple[int, str, bytes]):
+    # TODO: Infer the second argument as a type expression
+    # Should be fine
+    assert_type(a, tuple[int, str, bytes])  # error: [type-assertion-failure]
+
+    assert_type(a, tuple[int, str])  # error: [type-assertion-failure]
+    assert_type(a, tuple[int, str, bytes, None])  # error: [type-assertion-failure]
+    assert_type(a, tuple[int, bytes, str])  # error: [type-assertion-failure]
+
+def _(a: tuple[Any, ...]):
+    # TODO: Infer the second argument as a type expression
+    # Should be fine
+    assert_type(a, tuple[Any, ...])  # error: [type-assertion-failure]
+```
+
+## Unions
+
+Unions with the same elements are the same, regardless of order.
+
+```py
+from typing_extensions import assert_type
+
+def _(a: str | int):
+    # TODO: Infer the second argument as a type expression
+    # Should be fine
+    assert_type(a, str | int)  # error: [type-assertion-failure]
+```
+
+## Intersections
+
+Intersections are the same when their positive and negative parts are respectively the same,
+regardless of order.
+
+```py
+from typing_extensions import assert_type
+
+class A: ...
+class B: ...
+class C: ...
+class D: ...
+
+def _(a: A):
+    if isinstance(a, B) and not isinstance(a, C) and not isinstance(a, D):
+        reveal_type(a)  # revealed: A & B & ~C & ~D
+
+        # TODO: Use Python API to spell intersection type
+        # assert_type(a, B & A & ~D & ~C)
+```

crates/red_knot_python_semantic/src/types.rs

@@ -1,12 +1,14 @@
-use std::hash::Hash;
-
 use context::InferContext;
 use diagnostic::{report_not_iterable, report_not_iterable_possibly_unbound};
 use indexmap::IndexSet;
 use itertools::Itertools;
 use ruff_db::diagnostic::Severity;
 use ruff_db::files::File;
 use ruff_python_ast as ast;
+use rustc_hash::FxHasher;
+use std::collections::HashSet;
+use std::hash::{BuildHasherDefault, Hash};
+use std::iter;
 
 pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
 pub(crate) use self::diagnostic::register_lints;
@@ -547,6 +549,8 @@ pub enum Type<'db> {
     // TODO protocols, callable types, overloads, generics, type vars
 }
 
+type OrderedTypeSet<'a, 'db> = HashSet<&'a Type<'db>, BuildHasherDefault<FxHasher>>;
+
 impl<'db> Type<'db> {
     pub const fn is_unknown(&self) -> bool {
         matches!(self, Type::Unknown)
@@ -1408,6 +1412,89 @@ impl<'db> Type<'db> {
         }
     }
 
+    /// Returns true if this type and `other` are "exactly the same".
+    ///
+    /// This powers the `assert_type()` directive.
+    pub(crate) fn is_equals_to(self, db: &'db dyn Db, other: Type<'db>) -> bool {
+        let equal = |(first, second): (&Type<'db>, &Type<'db>)| first.is_equals_to(db, *second);
+
+        match (self, other) {
+            (Type::Todo(_), Type::Todo(_)) => false,
+
+            (Type::Any | Type::Unknown, Type::Any | Type::Unknown)
+            | (Type::Never, Type::Never)
+            | (Type::AlwaysTruthy, Type::AlwaysTruthy)
+            | (Type::AlwaysFalsy, Type::AlwaysFalsy)
+            | (Type::LiteralString, Type::LiteralString) => true,
+
+            (Type::KnownInstance(first), Type::KnownInstance(second)) => first == second,
+
+            (Type::FunctionLiteral(first), Type::FunctionLiteral(second)) => first == second,
+
+            (Type::ModuleLiteral(first), Type::ModuleLiteral(second)) => {
+                first.module(db) == second.module(db)
+            }
+
+            (Type::IntLiteral(first), Type::IntLiteral(second)) => first == second,
+            (Type::BooleanLiteral(first), Type::BooleanLiteral(second)) => first == second,
+            (Type::StringLiteral(first), Type::StringLiteral(second)) => {
+                first.value(db).as_ref() == second.value(db).as_ref()
+            }
+            (Type::BytesLiteral(first), Type::BytesLiteral(second)) => {
+                first.value(db).as_ref() == second.value(db).as_ref()
+            }
+            (Type::SliceLiteral(first), Type::SliceLiteral(second)) => {
+                first.as_tuple(db) == second.as_tuple(db)
+            }
+
+            (Type::ClassLiteral(first), Type::ClassLiteral(second)) => first.class == second.class,
+            (Type::Instance(first), Type::Instance(second)) => first.class == second.class,
+
+            (Type::SubclassOf(first), Type::SubclassOf(second)) => {
+                match (first.base, second.base) {
+                    (ClassBase::Todo(_), ClassBase::Todo(_)) => false,
+                    (ClassBase::Any | ClassBase::Unknown, ClassBase::Any | ClassBase::Unknown) => {
+                        true
+                    }
+                    (ClassBase::Class(first), ClassBase::Class(second)) => first == second,
+                    _ => false,
+                }
+            }
+
+            (Type::Tuple(first), Type::Tuple(second)) => {
+                first.len(db) == second.len(db)
+                    && iter::zip(first.elements(db), second.elements(db)).all(equal)
+            }
+
+            (Type::Union(first), Type::Union(second)) => {
+                let first_elements = first.elements(db);
+                let second_elements = second.elements(db);
+
+                if first_elements.len() != second_elements.len() {
+                    return false;
+                }
+
+                let first_elements = OrderedTypeSet::from_iter(first_elements.iter());
+                let second_elements = OrderedTypeSet::from_iter(second_elements.iter());
+
+                iter::zip(first_elements, second_elements).all(equal)
+            }
+
+            (Type::Intersection(first), Type::Intersection(second)) => {
+                let first_positive = OrderedTypeSet::from_iter(first.positive(db).iter());
+                let first_negative = OrderedTypeSet::from_iter(first.negative(db).iter());
+
+                let second_positive = OrderedTypeSet::from_iter(second.positive(db).iter());
+                let second_negative = OrderedTypeSet::from_iter(second.negative(db).iter());
+
+                iter::zip(first_positive, second_positive).all(equal)
+                    && iter::zip(first_negative, second_negative).all(equal)
+            }
+
+            _ => false,
+        }
+    }
+
     /// Return true if there is just a single inhabitant for this type.
     ///
     /// Note: This function aims to have no false positives, but might return `false`
@@ -1825,6 +1912,21 @@ impl<'db> Type<'db> {
                     CallOutcome::callable(len_ty.unwrap_or(normal_return_ty))
                 }
 
+                Some(KnownFunction::AssertType) => {
+                    let normal_return_ty = function_type.signature(db).return_ty;
+
+                    let [actual_type, asserted] = arg_types else {
+                        // TODO: Emit a diagnostic
+                        return CallOutcome::callable(normal_return_ty);
+                    };
+                    let Ok(asserted_type) = asserted.in_type_expression(db) else {
+                        // TODO: Emit a diagnostic
+                        return CallOutcome::callable(normal_return_ty);
+                    };
+
+                    CallOutcome::asserted(normal_return_ty, *actual_type, asserted_type)
+                }
+
                 _ => CallOutcome::callable(function_type.signature(db).return_ty),
             },
 
@@ -3086,13 +3188,15 @@ pub enum KnownFunction {
     Len,
     /// `typing(_extensions).final`
     Final,
+    /// `typing(_extensions).assert_type`
+    AssertType,
 }
 
 impl KnownFunction {
     pub fn constraint_function(self) -> Option<KnownConstraintFunction> {
         match self {
             Self::ConstraintFunction(f) => Some(f),
-            Self::RevealType | Self::Len | Self::Final => None,
+            Self::RevealType | Self::Len | Self::Final | Self::AssertType => None,
         }
     }
 
@@ -3111,6 +3215,7 @@ impl KnownFunction {
             ),
             "len" if definition.is_builtin_definition(db) => Some(KnownFunction::Len),
             "final" if definition.is_typing_definition(db) => Some(KnownFunction::Final),
+            "assert_type" if definition.is_typing_definition(db) => Some(KnownFunction::AssertType),
             _ => None,
         }
     }
@@ -4443,6 +4548,77 @@ pub(crate) mod tests {
         assert!(!from.into_type(&db).is_fully_static(&db));
     }
 
+    #[test_case(Ty::Any, Ty::Any)]
+    #[test_case(Ty::Unknown, Ty::Unknown)]
+    #[test_case(Ty::Any, Ty::Unknown)]
+    #[test_case(Ty::Never, Ty::Never)]
+    #[test_case(Ty::AlwaysTruthy, Ty::AlwaysTruthy)]
+    #[test_case(Ty::AlwaysFalsy, Ty::AlwaysFalsy)]
+    #[test_case(Ty::LiteralString, Ty::LiteralString)]
+    #[test_case(Ty::BooleanLiteral(true), Ty::BooleanLiteral(true))]
+    #[test_case(Ty::BooleanLiteral(false), Ty::BooleanLiteral(false))]
+    #[test_case(Ty::SliceLiteral(0, 1, 2), Ty::SliceLiteral(0, 1, 2))]
+    #[test_case(Ty::BuiltinClassLiteral("str"), Ty::BuiltinClassLiteral("str"))]
+    #[test_case(Ty::SubclassOfAny, Ty::SubclassOfUnknown)]
+    #[test_case(
+        Ty::Union(vec![Ty::BuiltinInstance("str"), Ty::BuiltinInstance("int")]),
+        Ty::Union(vec![Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str")])
+    )]
+    #[test_case(
+        Ty::Intersection {
+            pos: vec![Ty::BuiltinInstance("str"), Ty::BuiltinInstance("int")],
+            neg: vec![Ty::BuiltinInstance("bytes"), Ty::None]
+        },
+        Ty::Intersection {
+            pos: vec![Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str")],
+            neg: vec![Ty::None, Ty::BuiltinInstance("bytes")]
+        }
+    )]
+    #[test_case(
+        Ty::Intersection {
+            pos: vec![Ty::Union(vec![Ty::BuiltinInstance("str"), Ty::BuiltinInstance("int")])],
+            neg: vec![Ty::SubclassOfAny]
+        },
+        Ty::Intersection {
+            pos: vec![Ty::Union(vec![Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str")])],
+            neg: vec![Ty::SubclassOfUnknown]
+        }
+    )]
+    fn is_equals_to(a: Ty, b: Ty) {
+        let db = setup_db();
+        let a = a.into_type(&db);
+        let b = b.into_type(&db);
+
+        assert!(a.is_equals_to(&db, b));
+        assert!(b.is_equals_to(&db, a));
+    }
+
+    #[test_case(Ty::Todo, Ty::Todo)]
+    #[test_case(
+        Ty::Union(vec![Ty::BuiltinInstance("str"), Ty::BuiltinInstance("int")]),
+        Ty::Union(vec![Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str"), Ty::BuiltinInstance("bytes")])
+    )]
+    #[test_case(
+        Ty::Union(vec![Ty::BuiltinInstance("str"), Ty::BuiltinInstance("int"), Ty::BuiltinInstance("bytes")]),
+        Ty::Union(vec![Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str"), Ty::BuiltinInstance("dict")])
+    )]
+    #[test_case(
+        Ty::Tuple(vec![Ty::BuiltinInstance("str"), Ty::BuiltinInstance("int")]),
+        Ty::Tuple(vec![Ty::BuiltinInstance("str"), Ty::BuiltinInstance("int"), Ty::BuiltinInstance("bytes")])
+    )]
+    #[test_case(
+        Ty::Tuple(vec![Ty::BuiltinInstance("str"), Ty::BuiltinInstance("int")]),
+        Ty::Tuple(vec![Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str")])
+    )]
+    fn is_not_equal_to(a: Ty, b: Ty) {
+        let db = setup_db();
+        let a = a.into_type(&db);
+        let b = b.into_type(&db);
+
+        assert!(!a.is_equals_to(&db, b));
+        assert!(!b.is_equals_to(&db, a));
+    }
+
     #[test_case(Ty::IntLiteral(1); "is_int_literal_truthy")]
     #[test_case(Ty::IntLiteral(-1))]
     #[test_case(Ty::StringLiteral("foo"))]