Skip to content

Commit

Permalink
feat: correctly evaluate type of subscript
Browse files Browse the repository at this point in the history 
Signed-off-by: Shaygan <[email protected]>
  • Loading branch information
@Glyphack
Glyphack committed Aug 18, 2024
1 parent c06218f commit bebf2e0
Show file tree
Hide file tree
Showing 5 changed files with 114 additions and 73 deletions.
1 change: 0 additions & 1 deletion typechecker/src/build.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ use crate::{
},
settings::Settings,
symbol_table::{Id, SymbolTable},
types::PythonType,
};

#[derive(Debug)]
Expand Down
102 changes: 67 additions & 35 deletions typechecker/src/type_evaluator.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,7 @@ impl<'a> TypeEvaluator<'a> {
Ok(PythonType::Class(ClassType::new(
c.details.clone(),
vec![final_elm_type],
true,
)))
}
ast::Expression::Tuple(t) => {
Expand All @@ -170,6 +171,7 @@ impl<'a> TypeEvaluator<'a> {
Ok(PythonType::Class(ClassType::new(
c.details.clone(),
vec![elm_type],
true,
)))
}
ast::Expression::Dict(d) => {
Expand All @@ -184,6 +186,7 @@ impl<'a> TypeEvaluator<'a> {
Ok(PythonType::Class(ClassType::new(
c.details.clone(),
vec![key_type, value_type],
true,
)))
}
ast::Expression::Set(s) => {
Expand All @@ -198,6 +201,7 @@ impl<'a> TypeEvaluator<'a> {
Ok(PythonType::Class(ClassType::new(
class_type,
vec![elm_type],
true,
)))
}
ast::Expression::BoolOp(_) => Ok(self.get_builtin_type("bool").expect("typeshed")),
Expand Down Expand Up @@ -299,34 +303,9 @@ impl<'a> TypeEvaluator<'a> {
}
};
match value_type {
PythonType::Class(c) => {
let class_symbol_table_id = c.details.declaration_path.symbol_table_id;
let class_symbol_table = self
.imported_symbol_tables
.get(&class_symbol_table_id)
.unwrap();
let class_scope = c.details.class_scope_id;
let symbol_table_node =
class_symbol_table.lookup_attribute(&a.attr, class_scope);

match symbol_table_node {
Some(node) => self.get_symbol_type(node),
None => {
let scope = class_symbol_table
.get_scope_by_id(class_scope)
.unwrap_or_else(|| {
panic!(
"scope {:?} not found. available scopes are {:?}",
class_scope, symbol_table.scopes,
)
});
panic!(
"Attribute {:?} does not exist in scope {:?} on {:?}",
&a.attr, scope, c
);
}
}
}
PythonType::Class(c) => Ok(self
.lookup_on_class(symbol_table, &c, &a.attr)
.expect("attribute not found on type")),
PythonType::Module(module) => {
log::debug!("module: {:?}", module);
let module_sym_table =
Expand All @@ -342,9 +321,35 @@ impl<'a> TypeEvaluator<'a> {
&b.op,
)),
ast::Expression::Subscript(s) => {
let value_type = &self.get_type(&s.value, Some(symbol_table), None)?;
// This only handles container types and TODO
Ok(value_type.clone())
let value_type = self.get_type(&s.value, Some(symbol_table), None)?;

let typ = match value_type {
PythonType::Class(ref c) => {
if !c.is_instance {
return Ok(value_type);
}
let lookup_on_class = self.lookup_on_class(symbol_table, c, "__getitem__");
match lookup_on_class {
Some(PythonType::Callable(callable)) => {
let ret_type = self.get_return_type_of_callable(&callable, &[]);
match ret_type {
PythonType::TypeVar(ref tv) => {
let type_var_type = c.type_parameters.first().unwrap();
return Ok(type_var_type.clone());
}
_ => return Ok(ret_type),
}
}
_ => PythonType::Unknown,
}
}
_ => PythonType::Unknown,
};

if typ == PythonType::Unknown {
return Ok(value_type);
}
return Ok(typ);
}
ast::Expression::Slice(_) => Ok(PythonType::Unknown),
ast::Expression::Await(_) => Ok(PythonType::Unknown),
Expand Down Expand Up @@ -420,10 +425,11 @@ impl<'a> TypeEvaluator<'a> {
}
let type_parameters =
vec![self.get_annotation_type(&s.slice, symbol_table, None)];
PythonType::Class(ClassType {
details: class_type.details.clone(),
PythonType::Class(ClassType::new(
class_type.details.clone(),
type_parameters,
})
true,
))
}
// Illegal type annotation? Trying to subscript a non class type
Err(e) => PythonType::Unknown,
Expand Down Expand Up @@ -519,7 +525,11 @@ impl<'a> TypeEvaluator<'a> {
v.declaration_path.clone(),
decl_scope,
);
Ok(PythonType::Class(ClassType::new(class_symbol, vec![])))
Ok(PythonType::Class(ClassType::new(
class_symbol,
vec![],
false,
)))
} else {
Ok(var_type)
}
Expand Down Expand Up @@ -766,6 +776,7 @@ impl<'a> TypeEvaluator<'a> {
Ok(PythonType::Class(ClassType::new(
class_symbol.clone(),
class_def_type_parameters,
false,
)))
}
}
Expand Down Expand Up @@ -1071,4 +1082,25 @@ impl<'a> TypeEvaluator<'a> {
) -> PythonType {
f_type.return_type.clone()
}

fn lookup_on_class(
&self,
symbol_table: &SymbolTable,
c: &ClassType,
method_name: &str,
) -> Option<PythonType> {
dbg!(c);
let class_symbol_table_id = c.details.declaration_path.symbol_table_id;
let class_symbol_table = self
.imported_symbol_tables
.get(&class_symbol_table_id)
.unwrap();
let class_scope = c.details.class_scope_id;
let symbol_table_node = class_symbol_table.lookup_attribute(method_name, class_scope);

match symbol_table_node {
Some(node) => Some(self.get_symbol_type(node).expect("Cannot infer type")),
None => None,
}
}
}
12 changes: 11 additions & 1 deletion typechecker/src/types.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -115,13 +115,23 @@ pub struct ClassType {
pub details: symbol_table::Class,
// to represent types like `List[Int]`
pub type_parameters: Vec<PythonType>,
// Whether this class is an instance or the class itself.
// This is determined based on:
// 1. If the type is inferred from annotation of a parameter or variable that is an instance
// 2. If the type is inferred from a class node then it's an instance
pub is_instance: bool,
}

impl ClassType {
pub fn new(details: symbol_table::Class, type_parameters: Vec<PythonType>) -> Self {
pub fn new(
details: symbol_table::Class,
type_parameters: Vec<PythonType>,
is_instance: bool,
) -> Self {
Self {
details,
type_parameters,
is_instance,
}
}

Expand Down
Loading

0 comments on commit bebf2e0

Please sign in to comment.