Add decorator types to function type (#11253)

* Add `decorators: Vec<Type>` to `FunctionType` struct
* Thread decorators through two `add_function` definitions
* Populate decorators at the callsite in `infer_symbol_type`
* Small test
This commit is contained in:
plredmond 2024-05-02 16:58:56 -07:00 committed by GitHub
parent 59afff0e6a
commit b90a937a59
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 21 additions and 5 deletions

View file

@ -119,8 +119,9 @@ impl TypeStore {
self.modules.get(&file_id)
}
fn add_function(&self, file_id: FileId, name: &str) -> FunctionTypeId {
self.add_or_get_module(file_id).add_function(name)
fn add_function(&self, file_id: FileId, name: &str, decorators: Vec<Type>) -> FunctionTypeId {
self.add_or_get_module(file_id)
.add_function(name, decorators)
}
fn add_class(&self, file_id: FileId, name: &str, bases: Vec<Type>) -> ClassTypeId {
@ -306,9 +307,10 @@ impl ModuleTypeStore {
}
}
fn add_function(&mut self, name: &str) -> FunctionTypeId {
fn add_function(&mut self, name: &str, decorators: Vec<Type>) -> FunctionTypeId {
let func_id = self.functions.push(FunctionType {
name: Name::new(name),
decorators,
});
FunctionTypeId {
file_id: self.file_id,
@ -420,12 +422,17 @@ impl ClassType {
#[derive(Debug)]
pub(crate) struct FunctionType {
name: Name,
decorators: Vec<Type>,
}
impl FunctionType {
fn name(&self) -> &str {
self.name.as_str()
}
fn decorators(&self) -> &[Type] {
self.decorators.as_slice()
}
}
#[derive(Debug)]
@ -509,8 +516,9 @@ mod tests {
let store = TypeStore::default();
let files = Files::default();
let file_id = files.intern(Path::new("/foo"));
let id = store.add_function(file_id, "func");
let id = store.add_function(file_id, "func", vec![Type::Unknown]);
assert_eq!(store.get_function(id).name(), "func");
assert_eq!(store.get_function(id).decorators(), vec![Type::Unknown]);
let func = Type::Function(id);
assert_eq!(format!("{}", func.display(&store)), "func");
}

View file

@ -80,7 +80,15 @@ where
.resolve(ast.as_any_node_ref())
.expect("node key should resolve");
let ty = type_store.add_function(file_id, &node.name.id).into();
let decorator_tys = node
.decorator_list
.iter()
.map(|decorator| infer_expr_type(db, file_id, &decorator.expression))
.collect::<QueryResult<_>>()?;
let ty = type_store
.add_function(file_id, &node.name.id, decorator_tys)
.into();
type_store.cache_node_type(file_id, *node_key.erased(), ty);
ty
}