fold listcomp

This commit is contained in:
pca006132 2021-07-20 13:45:17 +08:00
parent 22455e43ac
commit fa31e8f336
1 changed files with 68 additions and 11 deletions

View File

@ -11,7 +11,7 @@ use itertools::izip;
use rustpython_parser::ast::{
self,
fold::{self, Fold},
Arguments, Expr, ExprKind, Located, Location,
Arguments, Comprehension, ExprKind, Located, Location,
};
pub struct PrimitiveStore {
@ -48,7 +48,9 @@ impl<'a> fold::Fold<()> for Inferencer<'a> {
ast::ExprKind::Lambda { args, body } => {
self.fold_lambda(node.location, *args, *body)?
}
ast::ExprKind::ListComp { elt, generators } => unimplemented!(),
ast::ExprKind::ListComp { elt, generators } => {
self.fold_listcomp(node.location, *elt, generators)?
}
_ => fold::fold_expr(self, node)?,
};
let custom = match &expr.node {
@ -70,18 +72,13 @@ impl<'a> fold::Fold<()> for Inferencer<'a> {
comparators,
} => Some(self.infer_compare(left, ops, comparators)?),
ast::ExprKind::Call { .. } => expr.custom,
ast::ExprKind::Subscript {
value,
slice,
ctx: _,
} => Some(self.infer_subscript(value.as_ref(), slice.as_ref())?),
ast::ExprKind::Subscript { value, slice, .. } => {
Some(self.infer_subscript(value.as_ref(), slice.as_ref())?)
}
ast::ExprKind::IfExp { test, body, orelse } => {
Some(self.infer_if_expr(test, body.as_ref(), orelse.as_ref())?)
}
ast::ExprKind::ListComp {
elt: _,
generators: _,
} => expr.custom, // already computed
ast::ExprKind::ListComp { .. } | ast::ExprKind::Lambda { .. } => expr.custom, // already computed
ast::ExprKind::Slice { .. } => None, // we don't need it for slice
_ => return Err("not supported yet".into()),
};
@ -186,6 +183,66 @@ impl<'a> Inferencer<'a> {
})
}
fn fold_listcomp(
&mut self,
location: Location,
elt: ast::Expr<()>,
mut generators: Vec<Comprehension>,
) -> Result<ast::Expr<Option<Type>>, String> {
if generators.len() != 1 {
return Err(
"Only 1 generator statement for list comprehension is supported.".to_string(),
);
}
let variable_mapping = self.variable_mapping.clone();
let mut new_context = Inferencer {
resolver: self.resolver,
unifier: self.unifier,
variable_mapping,
calls: self.calls,
primitives: self.primitives,
};
let elt = new_context.fold_expr(elt)?;
let generator = generators.pop().unwrap();
if generator.is_async {
return Err("Async iterator not supported.".to_string());
}
let target = new_context.fold_expr(*generator.target)?;
let iter = new_context.fold_expr(*generator.iter)?;
let ifs: Vec<_> = generator
.ifs
.into_iter()
.map(|v| new_context.fold_expr(v))
.collect::<Result<_, _>>()?;
// iter should be a list of targets...
// actually it should be an iterator of targets, but we don't have iter type for now
let list = new_context.unifier.add_ty(TypeEnum::TList {
ty: target.custom.unwrap(),
});
new_context.unifier.unify(iter.custom.unwrap(), list)?;
// if conditions should be bool
for v in ifs.iter() {
new_context
.unifier
.unify(v.custom.unwrap(), new_context.primitives.bool)?;
}
Ok(Located {
location,
custom: Some(list),
node: ExprKind::ListComp {
elt: Box::new(elt),
generators: vec![ast::Comprehension {
target: Box::new(target),
iter: Box::new(iter),
ifs,
is_async: false,
}],
},
})
}
fn infer_identifier(&mut self, id: &str) -> InferenceResult {
if let Some(ty) = self.variable_mapping.get(id) {
Ok(*ty)