diff --git a/nac3core/src/typecheck/type_inferencer/mod.rs b/nac3core/src/typecheck/type_inferencer/mod.rs index 23b9660e..624182b6 100644 --- a/nac3core/src/typecheck/type_inferencer/mod.rs +++ b/nac3core/src/typecheck/type_inferencer/mod.rs @@ -304,7 +304,9 @@ impl<'a> Inferencer<'a> { Ok(Located { location, - custom: Some(list), + custom: Some(new_context.unifier.add_ty(TypeEnum::TList { + ty: elt.custom.unwrap(), + })), node: ExprKind::ListComp { elt: Box::new(elt), generators: vec![ast::Comprehension { @@ -474,7 +476,7 @@ impl<'a> Inferencer<'a> { for t in elts.iter() { self.unifier.unify(ty, t.custom.unwrap())?; } - Ok(ty) + Ok(self.unifier.add_ty(TypeEnum::TList { ty })) } fn infer_tuple(&mut self, elts: &[ast::Expr>]) -> InferenceResult { diff --git a/nac3core/src/typecheck/type_inferencer/test.rs b/nac3core/src/typecheck/type_inferencer/test.rs index c06208a3..af937623 100644 --- a/nac3core/src/typecheck/type_inferencer/test.rs +++ b/nac3core/src/typecheck/type_inferencer/test.rs @@ -159,6 +159,13 @@ impl TestEnvironment { [("a", "fn[[x=bool], bool]"), ("b", "fn[[x=int32], int32]"), ("c", "bool"), ("d", "int32"), ("foo1", "Foo[bool]"), ("foo2", "Foo[int32]")].iter().cloned().collect() ; "obj test")] +#[test_case(indoc! {" + f = lambda x: True + a = [1, 2, 3] + b = [f(x) for x in a] + "}, + [("a", "list[int32]"), ("b", "list[bool]"), ("f", "fn[[x=int32], bool]")].iter().cloned().collect() + ; "listcomp test")] fn test_basic(source: &str, mapping: HashMap<&str, &str>) { let mut env = TestEnvironment::new(); let id_to_name = std::mem::take(&mut env.id_to_name);