[red-knot] Minor improvements to KnownFunction API (#15441)

A small PR to reduce some of the code duplication between the various
branches, make it a little more readable and move the API closer to what
we already have for `KnownClass`

crates/red_knot_python_semantic/src/module_resolver/module.rs

@@ -154,6 +154,10 @@ impl KnownModule {
         }
     }
 
+    pub const fn is_builtins(self) -> bool {
+        matches!(self, Self::Builtins)
+    }
+
     pub const fn is_typing(self) -> bool {
         matches!(self, Self::Typing)
     }

crates/red_knot_python_semantic/src/semantic_index/definition.rs

@@ -4,11 +4,10 @@ use ruff_python_ast as ast;
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::ast_node_ref::AstNodeRef;
-use crate::module_resolver::file_to_module;
 use crate::node_key::NodeKey;
 use crate::semantic_index::symbol::{FileScopeId, ScopeId, ScopedSymbolId};
 use crate::unpack::Unpack;
-use crate::{Db, KnownModule};
+use crate::Db;
 
 /// A definition of a symbol.
 ///
@@ -61,24 +60,6 @@ impl<'db> Definition<'db> {
     pub(crate) fn is_binding(self, db: &'db dyn Db) -> bool {
         self.kind(db).category().is_binding()
     }
-
-    pub(crate) fn is_builtin_definition(self, db: &'db dyn Db) -> bool {
-        file_to_module(db, self.file(db))
-            .is_some_and(|module| module.is_known(KnownModule::Builtins))
-    }
-
-    /// Return true if this symbol was defined in the `typing` or `typing_extensions` modules
-    pub(crate) fn is_typing_definition(self, db: &'db dyn Db) -> bool {
-        matches!(
-            file_to_module(db, self.file(db)).and_then(|module| module.known()),
-            Some(KnownModule::Typing | KnownModule::TypingExtensions)
-        )
-    }
-
-    pub(crate) fn is_knot_extensions_definition(self, db: &'db dyn Db) -> bool {
-        file_to_module(db, self.file(db))
-            .is_some_and(|module| module.is_known(KnownModule::KnotExtensions))
-    }
 }
 
 #[derive(Copy, Clone, Debug)]

crates/red_knot_python_semantic/src/types.rs

@@ -3401,47 +3401,51 @@ impl KnownFunction {
         definition: Definition<'db>,
         name: &str,
     ) -> Option<Self> {
-        match name {
+        let candidate = match name {
-            "reveal_type" if definition.is_typing_definition(db) => Some(KnownFunction::RevealType),
+            "isinstance" => Self::ConstraintFunction(KnownConstraintFunction::IsInstance),
-            "isinstance" if definition.is_builtin_definition(db) => Some(
+            "issubclass" => Self::ConstraintFunction(KnownConstraintFunction::IsSubclass),
-                KnownFunction::ConstraintFunction(KnownConstraintFunction::IsInstance),
+            "reveal_type" => Self::RevealType,
-            ),
+            "len" => Self::Len,
-            "issubclass" if definition.is_builtin_definition(db) => Some(
+            "final" => Self::Final,
-                KnownFunction::ConstraintFunction(KnownConstraintFunction::IsSubclass),
+            "no_type_check" => Self::NoTypeCheck,
-            ),
+            "assert_type" => Self::AssertType,
-            "len" if definition.is_builtin_definition(db) => Some(KnownFunction::Len),
+            "cast" => Self::Cast,
-            "final" if definition.is_typing_definition(db) => Some(KnownFunction::Final),
+            "static_assert" => Self::StaticAssert,
-            "no_type_check" if definition.is_typing_definition(db) => {
+            "is_subtype_of" => Self::IsSubtypeOf,
-                Some(KnownFunction::NoTypeCheck)
+            "is_disjoint_from" => Self::IsDisjointFrom,
-            }
+            "is_equivalent_to" => Self::IsEquivalentTo,
-            "assert_type" if definition.is_typing_definition(db) => Some(KnownFunction::AssertType),
+            "is_assignable_to" => Self::IsAssignableTo,
-            "cast" if definition.is_typing_definition(db) => Some(KnownFunction::Cast),
+            "is_fully_static" => Self::IsFullyStatic,
-            "static_assert" if definition.is_knot_extensions_definition(db) => {
+            "is_singleton" => Self::IsSingleton,
-                Some(KnownFunction::StaticAssert)
+            "is_single_valued" => Self::IsSingleValued,
-            }
+            _ => return None,
-            "is_subtype_of" if definition.is_knot_extensions_definition(db) => {
+        };
-                Some(KnownFunction::IsSubtypeOf)
+
-            }
+        candidate
-            "is_disjoint_from" if definition.is_knot_extensions_definition(db) => {
+            .check_module(file_to_module(db, definition.file(db))?.known()?)
-                Some(KnownFunction::IsDisjointFrom)
+            .then_some(candidate)
-            }
-            "is_equivalent_to" if definition.is_knot_extensions_definition(db) => {
-                Some(KnownFunction::IsEquivalentTo)
-            }
-            "is_assignable_to" if definition.is_knot_extensions_definition(db) => {
-                Some(KnownFunction::IsAssignableTo)
-            }
-            "is_fully_static" if definition.is_knot_extensions_definition(db) => {
-                Some(KnownFunction::IsFullyStatic)
-            }
-            "is_singleton" if definition.is_knot_extensions_definition(db) => {
-                Some(KnownFunction::IsSingleton)
-            }
-            "is_single_valued" if definition.is_knot_extensions_definition(db) => {
-                Some(KnownFunction::IsSingleValued)
     }
 
-            _ => None,
+    /// Return `true` if `self` is defined in `module` at runtime.
+    const fn check_module(self, module: KnownModule) -> bool {
+        match self {
+            Self::ConstraintFunction(constraint_function) => match constraint_function {
+                KnownConstraintFunction::IsInstance | KnownConstraintFunction::IsSubclass => {
+                    module.is_builtins()
+                }
+            },
+            Self::Len => module.is_builtins(),
+            Self::AssertType | Self::Cast | Self::RevealType | Self::Final | Self::NoTypeCheck => {
+                matches!(module, KnownModule::Typing | KnownModule::TypingExtensions)
+            }
+            Self::IsAssignableTo
+            | Self::IsDisjointFrom
+            | Self::IsEquivalentTo
+            | Self::IsFullyStatic
+            | Self::IsSingleValued
+            | Self::IsSingleton
+            | Self::IsSubtypeOf
+            | Self::StaticAssert => module.is_knot_extensions(),
         }
     }