[red-knot] Distribute intersections on negation

crates/red_knot_python_semantic/src/stdlib.rs

@@ -10,6 +10,7 @@ use crate::Db;
 enum CoreStdlibModule {
     Builtins,
     Types,
+    Typing,
     Typeshed,
     TypingExtensions,
 }
@@ -19,6 +20,7 @@ impl CoreStdlibModule {
         let module_name = match self {
             Self::Builtins => "builtins",
             Self::Types => "types",
+            Self::Typing => "typing",
             Self::Typeshed => "_typeshed",
             Self::TypingExtensions => "typing_extensions",
         };
@@ -63,6 +65,13 @@ pub(crate) fn types_symbol_ty<'db>(db: &'db dyn Db, symbol: &str) -> Type<'db> {
     core_module_symbol_ty(db, CoreStdlibModule::Types, symbol)
 }
 
+/// Lookup the type of `symbol` in the `typing` module namespace.
+///
+/// Returns `Unbound` if the `types` module isn't available for some reason.
+#[inline]
+pub(crate) fn typing_symbol_ty<'db>(db: &'db dyn Db, symbol: &str) -> Type<'db> {
+    core_module_symbol_ty(db, CoreStdlibModule::Typing, symbol)
+}
 /// Lookup the type of `symbol` in the `_typeshed` module namespace.
 ///
 /// Returns `Unbound` if the `_typeshed` module isn't available for some reason.

crates/red_knot_python_semantic/src/types.rs

@@ -12,6 +12,7 @@ use crate::semantic_index::{
 };
 use crate::stdlib::{
     builtins_symbol_ty, types_symbol_ty, typeshed_symbol_ty, typing_extensions_symbol_ty,
+    typing_symbol_ty,
 };
 use crate::types::narrow::narrowing_constraint;
 use crate::{Db, FxOrderSet, HasTy, Module, SemanticModel};
@@ -740,7 +741,9 @@ impl<'db> Type<'db> {
                     | KnownClass::Dict
                     | KnownClass::GenericAlias
                     | KnownClass::ModuleType
-                    | KnownClass::FunctionType,
+                    | KnownClass::FunctionType
+                    | KnownClass::Sized
+                    | KnownClass::Hashable,
                 ) => false,
                 None => false,
             },
@@ -1132,6 +1135,9 @@ pub enum KnownClass {
     GenericAlias,
     ModuleType,
     FunctionType,
+    // Typing
+    Sized,
+    Hashable,
     // Typeshed
     NoneType, // Part of `types` for Python >= 3.10
 }
@@ -1153,6 +1159,8 @@ impl<'db> KnownClass {
             Self::GenericAlias => "GenericAlias",
             Self::ModuleType => "ModuleType",
             Self::FunctionType => "FunctionType",
+            Self::Sized => "Sized",
+            Self::Hashable => "Hashable",
             Self::NoneType => "NoneType",
         }
     }
@@ -1177,6 +1185,8 @@ impl<'db> KnownClass {
             Self::GenericAlias | Self::ModuleType | Self::FunctionType => {
                 types_symbol_ty(db, self.as_str())
             }
+            Self::Sized | Self::Hashable => typing_symbol_ty(db, self.as_str()),
+
             Self::NoneType => typeshed_symbol_ty(db, self.as_str()),
         }
     }
@@ -1210,6 +1220,8 @@ impl<'db> KnownClass {
             "NoneType" => Some(Self::NoneType),
             "ModuleType" => Some(Self::ModuleType),
             "FunctionType" => Some(Self::FunctionType),
+            "Sized" => Some(Self::Sized),
+            "Hashable" => Some(Self::Hashable),
             _ => None,
         }
     }
@@ -1232,6 +1244,9 @@ impl<'db> KnownClass {
             | Self::Set
             | Self::Dict => module.name() == "builtins",
             Self::GenericAlias | Self::ModuleType | Self::FunctionType => module.name() == "types",
+            Self::Sized | Self::Hashable => {
+                matches!(module.name().as_str(), "typing" | "collections.abc")
+            }
             Self::NoneType => matches!(module.name().as_str(), "_typeshed" | "types"),
         }
     }

crates/red_knot_python_semantic/src/types/builder.rs

@@ -177,6 +177,33 @@ impl<'db> IntersectionBuilder<'db> {
                 self = self.add_negative(*elem);
             }
             self
+        } else if let Type::Intersection(intersection) = ty {
+            // (A | B) & !(C & !D) -> (A | B) & (!C | D)
+            // -> ((A | B) & !C) | ((A | B) & D)
+            //
+            // i.e. if we have an intersection of positive constraints C
+            // and negative constraints D, then our new intersection
+            // is (existing & !C) | (existing & D)
+
+            let positive_side = intersection
+                .positive(self.db)
+                .iter()
+                // we negate all the positive constraints while distributing
+                .map(|elem| self.clone().add_negative(*elem));
+
+            let negative_side = intersection
+                .negative(self.db)
+                .iter()
+                // all negative constraints end up becoming positive constraints
+                .map(|elem| self.clone().add_positive(*elem));
+
+            positive_side.chain(negative_side).fold(
+                IntersectionBuilder::empty(self.db),
+                |mut builder, sub| {
+                    builder.intersections.extend(sub.intersections);
+                    builder
+                },
+            )
         } else {
             for inner in &mut self.intersections {
                 inner.add_negative(self.db, ty);
@@ -368,10 +395,17 @@ mod tests {
     use crate::program::{Program, SearchPathSettings};
     use crate::python_version::PythonVersion;
     use crate::types::{KnownClass, StringLiteralType, UnionBuilder};
+    use crate::Db;
     use crate::ProgramSettings;
     use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
     use test_case::test_case;
 
+    impl<'db> IntersectionBuilder<'db> {
+        fn negation_of(db: &'db dyn Db, ty: Type<'db>) -> Type<'db> {
+            IntersectionBuilder::new(db).add_negative(ty).build()
+        }
+    }
+
     fn setup_db() -> TestDb {
         let db = TestDb::new();
 
@@ -591,6 +625,32 @@ mod tests {
         assert_eq!(i1.pos_vec(&db), &[ta, t1]);
     }
 
+    #[test]
+    fn intersection_negation_distributes_over_union() {
+        let db = setup_db();
+        let st = KnownClass::Sized.to_instance(&db);
+        let ht = KnownClass::Hashable.to_instance(&db);
+        // sh_t: Sized & Hashable
+        let sh_t = IntersectionBuilder::new(&db)
+            .add_positive(st)
+            .add_positive(ht)
+            .build()
+            .expect_intersection();
+        assert_eq!(sh_t.pos_vec(&db), &[st, ht]);
+        assert_eq!(sh_t.neg_vec(&db), &[]);
+
+        // !sh_t => not Sized or not Hashable
+        let not_s_h_t = IntersectionBuilder::new(&db)
+            .add_negative(Type::Intersection(sh_t))
+            .build()
+            .expect_union();
+
+        // should have (not Sized) or (not Hashable)
+        let not_st = IntersectionBuilder::negation_of(&db, st);
+        let not_ht = IntersectionBuilder::negation_of(&db, ht);
+
+        assert_eq!(not_s_h_t.elements(&db), &[not_st, not_ht]);
+    }
     #[test]
     fn build_intersection_self_negation() {
         let db = setup_db();