[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`
This commit is contained in:
Alex Waygood 2025-01-12 16:06:31 +00:00 committed by GitHub
parent c8795fcb37
commit 06b7f4495e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 49 additions and 60 deletions

View file

@ -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)
}

View file

@ -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)]

View file

@ -3401,47 +3401,51 @@ impl KnownFunction {
definition: Definition<'db>,
name: &str,
) -> Option<Self> {
match name {
"reveal_type" if definition.is_typing_definition(db) => Some(KnownFunction::RevealType),
"isinstance" if definition.is_builtin_definition(db) => Some(
KnownFunction::ConstraintFunction(KnownConstraintFunction::IsInstance),
),
"issubclass" if definition.is_builtin_definition(db) => Some(
KnownFunction::ConstraintFunction(KnownConstraintFunction::IsSubclass),
),
"len" if definition.is_builtin_definition(db) => Some(KnownFunction::Len),
"final" if definition.is_typing_definition(db) => Some(KnownFunction::Final),
"no_type_check" if definition.is_typing_definition(db) => {
Some(KnownFunction::NoTypeCheck)
}
"assert_type" if definition.is_typing_definition(db) => Some(KnownFunction::AssertType),
"cast" if definition.is_typing_definition(db) => Some(KnownFunction::Cast),
"static_assert" if definition.is_knot_extensions_definition(db) => {
Some(KnownFunction::StaticAssert)
}
"is_subtype_of" if definition.is_knot_extensions_definition(db) => {
Some(KnownFunction::IsSubtypeOf)
}
"is_disjoint_from" if definition.is_knot_extensions_definition(db) => {
Some(KnownFunction::IsDisjointFrom)
}
"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)
let candidate = match name {
"isinstance" => Self::ConstraintFunction(KnownConstraintFunction::IsInstance),
"issubclass" => Self::ConstraintFunction(KnownConstraintFunction::IsSubclass),
"reveal_type" => Self::RevealType,
"len" => Self::Len,
"final" => Self::Final,
"no_type_check" => Self::NoTypeCheck,
"assert_type" => Self::AssertType,
"cast" => Self::Cast,
"static_assert" => Self::StaticAssert,
"is_subtype_of" => Self::IsSubtypeOf,
"is_disjoint_from" => Self::IsDisjointFrom,
"is_equivalent_to" => Self::IsEquivalentTo,
"is_assignable_to" => Self::IsAssignableTo,
"is_fully_static" => Self::IsFullyStatic,
"is_singleton" => Self::IsSingleton,
"is_single_valued" => Self::IsSingleValued,
_ => return None,
};
candidate
.check_module(file_to_module(db, definition.file(db))?.known()?)
.then_some(candidate)
}
_ => 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(),
}
}