[red-knot] simplify subtypes from unions (#13401)

Add `Type::is_subtype_of` method, and simplify subtypes out of unions.
astral-sh · Sep 19, 2024 · cf1e91b · cf1e91b
1 parent 125eaaf
commit cf1e91b
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 9 deletions.
diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
@@ -388,16 +388,18 @@ impl<'db> Type<'db> {
         }
     }
 
-    /// Return true if this type is [assignable to] type `target`.
+    /// Return true if this type is a [subtype of] type `target`.
     ///
-    /// [assignable to]: https://typing.readthedocs.io/en/latest/spec/concepts.html#the-assignable-to-or-consistent-subtyping-relation
-    pub(crate) fn is_assignable_to(self, db: &'db dyn Db, target: Type<'db>) -> bool {
+    /// [subtype of]: https://typing.readthedocs.io/en/latest/spec/concepts.html#subtype-supertype-and-type-equivalence
+    pub(crate) fn is_subtype_of(self, db: &'db dyn Db, target: Type<'db>) -> bool {
         if self.is_equivalent_to(db, target) {
             return true;
         }
         match (self, target) {
-            (Type::Unknown | Type::Any | Type::Never, _) => true,
-            (_, Type::Unknown | Type::Any) => true,
+            (Type::Unknown | Type::Any, _) => false,
+            (_, Type::Unknown | Type::Any) => false,
+            (Type::Never, _) => true,
+            (_, Type::Never) => false,
             (Type::IntLiteral(_), Type::Instance(class))
                 if class.is_stdlib_symbol(db, "builtins", "int") =>
             {
@@ -417,12 +419,28 @@ impl<'db> Type<'db> {
             (ty, Type::Union(union)) => union
                 .elements(db)
                 .iter()
-                .any(|&elem_ty| ty.is_assignable_to(db, elem_ty)),
+                .any(|&elem_ty| ty.is_subtype_of(db, elem_ty)),
             // TODO
             _ => false,
         }
     }
 
+    /// Return true if this type is [assignable to] type `target`.
+    ///
+    /// [assignable to]: https://typing.readthedocs.io/en/latest/spec/concepts.html#the-assignable-to-or-consistent-subtyping-relation
+    pub(crate) fn is_assignable_to(self, db: &'db dyn Db, target: Type<'db>) -> bool {
+        match (self, target) {
+            (Type::Unknown | Type::Any, _) => true,
+            (_, Type::Unknown | Type::Any) => true,
+            (ty, Type::Union(union)) => union
+                .elements(db)
+                .iter()
+                .any(|&elem_ty| ty.is_assignable_to(db, elem_ty)),
+            // TODO other types containing gradual forms (e.g. generics containing Any/Unknown)
+            _ => self.is_subtype_of(db, target),
+        }
+    }
+
     /// Return true if this type is equivalent to type `other`.
     pub(crate) fn is_equivalent_to(self, _db: &'db dyn Db, other: Type<'db>) -> bool {
         // TODO equivalent but not identical structural types, differently-ordered unions and
@@ -1132,6 +1150,31 @@ mod tests {
         assert!(!from.into_type(&db).is_assignable_to(&db, to.into_type(&db)));
     }
 
+    #[test_case(Ty::Never, Ty::IntLiteral(1))]
+    #[test_case(Ty::IntLiteral(1), Ty::BuiltinInstance("int"))]
+    #[test_case(Ty::StringLiteral("foo"), Ty::BuiltinInstance("str"))]
+    #[test_case(Ty::StringLiteral("foo"), Ty::LiteralString)]
+    #[test_case(Ty::LiteralString, Ty::BuiltinInstance("str"))]
+    #[test_case(Ty::BytesLiteral("foo"), Ty::BuiltinInstance("bytes"))]
+    #[test_case(Ty::IntLiteral(1), Ty::Union(vec![Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str")]))]
+    fn is_subtype_of(from: Ty, to: Ty) {
+        let db = setup_db();
+        assert!(from.into_type(&db).is_subtype_of(&db, to.into_type(&db)));
+    }
+
+    #[test_case(Ty::Unknown, Ty::IntLiteral(1))]
+    #[test_case(Ty::Any, Ty::IntLiteral(1))]
+    #[test_case(Ty::IntLiteral(1), Ty::Unknown)]
+    #[test_case(Ty::IntLiteral(1), Ty::Any)]
+    #[test_case(Ty::IntLiteral(1), Ty::Union(vec![Ty::Unknown, Ty::BuiltinInstance("str")]))]
+    #[test_case(Ty::IntLiteral(1), Ty::BuiltinInstance("str"))]
+    #[test_case(Ty::BuiltinInstance("int"), Ty::BuiltinInstance("str"))]
+    #[test_case(Ty::BuiltinInstance("int"), Ty::IntLiteral(1))]
+    fn is_not_subtype_of(from: Ty, to: Ty) {
+        let db = setup_db();
+        assert!(!from.into_type(&db).is_subtype_of(&db, to.into_type(&db)));
+    }
+
     #[test_case(
         Ty::Union(vec![Ty::IntLiteral(1), Ty::IntLiteral(2)]),
         Ty::Union(vec![Ty::IntLiteral(1), Ty::IntLiteral(2)])

diff --git a/crates/red_knot_python_semantic/src/types/builder.rs b/crates/red_knot_python_semantic/src/types/builder.rs
@@ -46,10 +46,23 @@ impl<'db> UnionBuilder<'db> {
     pub(crate) fn add(mut self, ty: Type<'db>) -> Self {
         match ty {
             Type::Union(union) => {
-                self.elements.extend(union.elements(self.db));
+                for element in union.elements(self.db) {
+                    self = self.add(*element);
+                }
             }
             Type::Never => {}
             _ => {
+                let mut remove = vec![];
+                for element in &self.elements {
+                    if ty.is_subtype_of(self.db, *element) {
+                        return self;
+                    } else if element.is_subtype_of(self.db, ty) {
+                        remove.push(*element);
+                    }
+                }
+                for element in remove {
+                    self.elements.remove(&element);
+                }
                 self.elements.insert(ty);
             }
         }
@@ -368,6 +381,24 @@ mod tests {
         assert_eq!(union.elements_vec(&db), &[t0, t1, t2]);
     }
 
+    #[test]
+    fn build_union_simplify_subtype() {
+        let db = setup_db();
+        let t0 = builtins_symbol_ty(&db, "str").to_instance(&db);
+        let t1 = Type::LiteralString;
+        let t2 = Type::Unknown;
+        let u0 = UnionType::from_elements(&db, [t0, t1]);
+        let u1 = UnionType::from_elements(&db, [t1, t0]);
+        let u2 = UnionType::from_elements(&db, [t0, t1, t2]);
+
+        assert_eq!(u0, t0);
+        assert_eq!(u1, t0);
+        assert_eq!(u2.expect_union().elements_vec(&db), &[t0, t2]);
+    }
+
+    #[test]
+    fn build_union_no_simplify_any() {}
+
     impl<'db> IntersectionType<'db> {
         fn pos_vec(self, db: &'db TestDb) -> Vec<Type<'db>> {
             self.positive(db).into_iter().copied().collect()

diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -5800,8 +5800,7 @@ mod tests {
         .unwrap();
         db.write_file("/src/c.pyi", "x: int").unwrap();
 
-        // TODO this should simplify to just 'int'
-        assert_public_ty(&db, "/src/a.py", "x", "int | Literal[1]");
+        assert_public_ty(&db, "/src/a.py", "x", "int");
     }
 
     // Incremental inference tests