fix: collection bugs

Cargo.lock

@@ -150,9 +150,9 @@ checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0"
 
 [[package]]
 name = "els"
-version = "0.1.63-nightly.1"
+version = "0.1.63-nightly.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "499bf0ea920c33a3b8f905450435135a8fbefc6c06c3b53f2b5c0e66d25367b5"
+checksum = "378a27662786b3f32abd006ebcda0f1c602f77edef1ec7739f01a15869a5f6c7"
 dependencies = [
  "erg_common",
  "erg_compiler",
@@ -166,9 +166,9 @@ dependencies = [
 
 [[package]]
 name = "erg_common"
-version = "0.6.51-nightly.1"
+version = "0.6.51-nightly.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1d112c6737def3a5ee868510f4ce9ef9be64b2a0bdf8899b24576a1e8521d401"
+checksum = "8ad4be8c3b49a7907528ceba93cb4cf9a1e5fc12abe9ebc78e325b610fc13162"
 dependencies = [
  "backtrace-on-stack-overflow",
  "erg_proc_macros",
@@ -179,9 +179,9 @@ dependencies = [
 
 [[package]]
 name = "erg_compiler"
-version = "0.6.51-nightly.1"
+version = "0.6.51-nightly.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "32d241b882842c86c6c3e2aa54c4c8121fdbedd8a15a7b563cf5f7c1c5c0c0bb"
+checksum = "cf949087e14c3c1b927c32dcf704c90b4f411c52954dc66321bc60d2f8392061"
 dependencies = [
  "erg_common",
  "erg_parser",
@@ -189,9 +189,9 @@ dependencies = [
 
 [[package]]
 name = "erg_parser"
-version = "0.6.51-nightly.1"
+version = "0.6.51-nightly.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "91efa61b1495ee680394b6f9d98d84535dc8445329f0310734a3e40671a0061a"
+checksum = "a720002b515aa3a0818fbce89ce32bb23b3ee372f143640769d43e5722c82453"
 dependencies = [
  "erg_common",
  "erg_proc_macros",
@@ -200,9 +200,9 @@ dependencies = [
 
 [[package]]
 name = "erg_proc_macros"
-version = "0.6.51-nightly.1"
+version = "0.6.51-nightly.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bd78027648c82db68efed1a1789bd84ce09f07ea7630085c9015c4ed6b1a192d"
+checksum = "c359be27ae21d38e9350e70eb3c4156737aa4e5ba7e59ab3b439a90167b59c97"
 dependencies = [
  "quote",
  "syn 1.0.109",

Cargo.toml

@@ -24,9 +24,9 @@ edition = "2021"
 repository = "https://github.com/mtshiba/pylyzer"
 
 [workspace.dependencies]
-erg_common = { version = "0.6.51-nightly.1", features = ["py_compat", "els"] }
+erg_common = { version = "0.6.51-nightly.3", features = ["py_compat", "els"] }
-erg_compiler = { version = "0.6.51-nightly.1", features = ["py_compat", "els"] }
+erg_compiler = { version = "0.6.51-nightly.3", features = ["py_compat", "els"] }
-els = { version = "0.1.63-nightly.1", features = ["py_compat"] }
+els = { version = "0.1.63-nightly.3", features = ["py_compat"] }
 # rustpython-parser = { version = "0.3.0", features = ["all-nodes-with-ranges", "location"] }
 # rustpython-ast = { version = "0.3.0", features = ["all-nodes-with-ranges", "location"] }
 rustpython-parser = { git = "https://github.com/RustPython/Parser", version = "0.4.0", features = ["all-nodes-with-ranges", "location"] }

crates/py2erg/convert.rs

@@ -13,14 +13,14 @@ use erg_compiler::erg_parser::ast::{
     Accessor, Args, BinOp, Block, ClassAttr, ClassAttrs, ClassDef, Compound, ConstAccessor,
     ConstApp, ConstArgs, ConstAttribute, ConstBinOp, ConstBlock, ConstDict, ConstExpr,
     ConstKeyValue, ConstLambda, ConstList, ConstListWithLength, ConstNormalList, ConstNormalSet,
-    ConstPosArg, ConstSet, Decorator, Def, DefBody, DefId, DefaultParamSignature, Dict, Dummy,
+    ConstPosArg, ConstSet, Decorator, Def, DefBody, DefId, DefaultParamSignature, Dict,
-    Expr, Identifier, KeyValue, KwArg, Lambda, LambdaSignature, List, ListComprehension, Literal,
+    DictComprehension, Dummy, Expr, Identifier, KeyValue, KwArg, Lambda, LambdaSignature, List,
-    Methods, Module, NonDefaultParamSignature, NormalDict, NormalList, NormalRecord, NormalSet,
+    ListComprehension, Literal, Methods, Module, NonDefaultParamSignature, NormalDict, NormalList,
-    NormalTuple, ParamPattern, ParamTySpec, Params, PosArg, PreDeclTypeSpec, ReDef, Record,
+    NormalRecord, NormalSet, NormalTuple, ParamPattern, ParamTySpec, Params, PosArg,
-    RecordAttrs, Set, SetComprehension, Signature, SubrSignature, SubrTypeSpec, Tuple,
+    PreDeclTypeSpec, ReDef, Record, RecordAttrs, Set, SetComprehension, Signature, SubrSignature,
-    TupleTypeSpec, TypeAppArgs, TypeAppArgsKind, TypeAscription, TypeBoundSpec, TypeBoundSpecs,
+    SubrTypeSpec, Tuple, TupleTypeSpec, TypeAppArgs, TypeAppArgsKind, TypeAscription,
-    TypeSpec, TypeSpecWithOp, UnaryOp, VarName, VarPattern, VarRecordAttr, VarRecordAttrs,
+    TypeBoundSpec, TypeBoundSpecs, TypeSpec, TypeSpecWithOp, UnaryOp, VarName, VarPattern,
-    VarRecordPattern, VarSignature, VisModifierSpec,
+    VarRecordAttr, VarRecordAttrs, VarRecordPattern, VarSignature, VisModifierSpec,
 };
 use erg_compiler::erg_parser::desugar::Desugarer;
 use erg_compiler::erg_parser::token::{Token, TokenKind, AS, COLON, DOT, EQUAL};
@@ -1017,48 +1017,50 @@ impl ASTConverter {
 
     fn convert_ident_type_spec(&mut self, name: String, range: PySourceRange) -> TypeSpec {
         let loc = pyloc_to_ergloc(range);
+        let global = ConstExpr::Accessor(ConstAccessor::Local(Identifier::private_with_loc(
+            "global".into(),
+            loc,
+        )));
+        let obj = ConstExpr::Accessor(ConstAccessor::Local(Identifier::private_with_loc(
+            "Obj".into(),
+            loc,
+        )));
         match &name[..] {
+            "dict" => {
+                let kv = ConstKeyValue::new(obj.clone(), obj.clone());
+                let (l, r) = Self::gen_enclosure_tokens(TokenKind::LSqBr, range);
+                let dict = ConstDict::new(l, r, vec![kv]);
+                TypeSpec::poly(
+                    global.attr(Identifier::private_with_loc("Dict!".into(), loc)),
+                    ConstArgs::single(ConstExpr::Dict(dict)),
+                )
+            }
+            "set" => TypeSpec::poly(
+                global.attr(Identifier::private_with_loc("Set!".into(), loc)),
+                ConstArgs::single(obj),
+            ),
+            "tuple" => TypeSpec::poly(
+                global.attr(Identifier::private_with_loc("HomogenousTuple".into(), loc)),
+                ConstArgs::single(obj),
+            ),
             // Iterable[T] => Iterable(T), Iterable => Iterable(Obj)
             global_unary_collections!() => TypeSpec::poly(
-                ConstExpr::Accessor(ConstAccessor::Local(Identifier::private_with_loc(
+                global.attr(Identifier::private_with_loc(name.into(), loc)),
-                    "global".into(),
+                ConstArgs::single(obj),
-                    loc,
-                )))
-                .attr(Identifier::private_with_loc(name.into(), loc)),
-                ConstArgs::single(ConstExpr::Accessor(ConstAccessor::Local(
-                    Identifier::private_with_loc("Obj".into(), loc),
-                ))),
             ),
             // MutableSequence[T] => Sequence!(T), MutableSequence => Sequence!(Obj)
             global_mutable_unary_collections!() => TypeSpec::poly(
-                ConstExpr::Accessor(ConstAccessor::Local(Identifier::private_with_loc(
+                global.attr(Identifier::private_with_loc(
-                    "global".into(),
-                    loc,
-                )))
-                .attr(Identifier::private_with_loc(
                     format!("{}!", name.trim_start_matches("Mutable")).into(),
                     loc,
                 )),
-                ConstArgs::single(ConstExpr::Accessor(ConstAccessor::Local(
+                ConstArgs::single(obj),
-                    Identifier::private_with_loc("Obj".into(), loc),
-                ))),
             ),
             // Mapping => Mapping(Obj, Obj)
             global_binary_collections!() => TypeSpec::poly(
-                ConstExpr::Accessor(ConstAccessor::Local(Identifier::private_with_loc(
+                global.attr(Identifier::private_with_loc(name.into(), loc)),
-                    "global".into(),
-                    loc,
-                )))
-                .attr(Identifier::private_with_loc(name.into(), loc)),
                 ConstArgs::pos_only(
-                    vec![
+                    vec![ConstPosArg::new(obj.clone()), ConstPosArg::new(obj)],
-                        ConstPosArg::new(ConstExpr::Accessor(ConstAccessor::Local(
-                            Identifier::private_with_loc("Obj".into(), loc),
-                        ))),
-                        ConstPosArg::new(ConstExpr::Accessor(ConstAccessor::Local(
-                            Identifier::private_with_loc("Obj".into(), loc),
-                        ))),
-                    ],
                     None,
                 ),
             ),
@@ -1502,6 +1504,22 @@ impl ASTConverter {
                 let tuple = TupleTypeSpec::new(Some((l.loc(), r.loc())), tys);
                 TypeSpec::Tuple(tuple)
             }
+            "Set" | "set" => {
+                let elem_t = match self.convert_expr_to_const(args) {
+                    Some(elem_t) => elem_t,
+                    None => {
+                        ConstExpr::Accessor(ConstAccessor::Local(Identifier::private("Obj".into())))
+                    }
+                };
+                let elem_t = ConstPosArg::new(elem_t);
+                let global =
+                    ConstExpr::Accessor(ConstAccessor::Local(Identifier::private("global".into())));
+                let acc = ConstAccessor::Attr(ConstAttribute::new(
+                    global,
+                    Identifier::private("Set!".into()),
+                ));
+                TypeSpec::poly(acc, ConstArgs::pos_only(vec![elem_t], None))
+            }
             _ => Self::gen_dummy_type_spec(args.location()),
         }
     }
@@ -1899,8 +1917,8 @@ impl ASTConverter {
                     .map(|ex| PosArg::new(self.convert_expr(ex)))
                     .collect::<Vec<_>>();
                 let elems = Args::pos_only(elements, None);
-                Expr::Set(Set::Normal(NormalSet::new(l_brace, r_brace, elems)))
+                let set = Expr::Set(Set::Normal(NormalSet::new(l_brace, r_brace, elems)));
-                // Self::mutate_expr(set)
+                Self::mutate_expr(set)
             }
             py_ast::Expr::SetComp(comp) => {
                 let (l_brace, r_brace) = Self::gen_enclosure_tokens(TokenKind::LBrace, comp.range);
@@ -1919,14 +1937,14 @@ impl ASTConverter {
                     .next()
                     .map(|ex| self.convert_expr(ex));
                 let generators = vec![(ident, iter)];
-                Expr::Set(Set::Comprehension(SetComprehension::new(
+                let set = Expr::Set(Set::Comprehension(SetComprehension::new(
                     l_brace,
                     r_brace,
                     Some(layout),
                     generators,
                     guard,
-                )))
+                )));
-                // Self::mutate_expr(set)
+                Self::mutate_expr(set)
             }
             py_ast::Expr::Dict(dict) => {
                 let (l_brace, r_brace) = Self::gen_enclosure_tokens(TokenKind::LBrace, dict.range);
@@ -1945,6 +1963,30 @@ impl ASTConverter {
                 let dict = Expr::Dict(Dict::Normal(NormalDict::new(l_brace, r_brace, kvs)));
                 Self::mutate_expr(dict)
             }
+            py_ast::Expr::DictComp(comp) => {
+                let (l_brace, r_brace) = Self::gen_enclosure_tokens(TokenKind::LBrace, comp.range);
+                let key = self.convert_expr(*comp.key);
+                let value = self.convert_expr(*comp.value);
+                let kv = KeyValue::new(key, value);
+                let generator = comp.generators.into_iter().next().unwrap();
+                let target = self.convert_expr(generator.target);
+                let Expr::Accessor(Accessor::Ident(ident)) = target else {
+                    log!(err "unimplemented: {target}");
+                    let loc = pyloc_to_ergloc(comp.range);
+                    return Expr::Dummy(Dummy::new(Some(loc), vec![]));
+                };
+                let iter = self.convert_expr(generator.iter);
+                let guard = generator
+                    .ifs
+                    .into_iter()
+                    .next()
+                    .map(|ex| self.convert_expr(ex));
+                let generators = vec![(ident, iter)];
+                let dict = Expr::Dict(Dict::Comprehension(DictComprehension::new(
+                    l_brace, r_brace, kv, generators, guard,
+                )));
+                Self::mutate_expr(dict)
+            }
             py_ast::Expr::Tuple(tuple) => {
                 let (l, r) = Self::gen_enclosure_tokens(TokenKind::LParen, tuple.range);
                 let elements = tuple

crates/py2erg/gen_decl.rs

@@ -102,7 +102,7 @@ impl DeclFileGenerator {
                 let typ = self.escape_type(ref_t.replace_failure().to_string_unabbreviated());
                 // Erg can automatically import nested modules
                 // `import http.client` => `http = pyimport "http"`
-                let decl = if ref_t.is_py_module() {
+                let decl = if ref_t.is_py_module() && ref_t.typarams()[0].is_str_value() {
                     name = name.split('.').next().unwrap().to_string();
                     let full_path_str = ref_t.typarams()[0].to_string_unabbreviated();
                     let mod_name = mod_name(Path::new(full_path_str.trim_matches('"')));

tests/collection.py

@@ -23,6 +23,10 @@ _ = t[0] == 1 # OK
 _ = t[1] == 1 # ERR
 _ = t[0:1]
 
+s: set[int] = {1, 2}
+s.add(1)
+s.add("a") # ERR
+
 def f(s: Str): return None
 for i in getattr(1, "aaa", ()):
     f(i)
@@ -32,3 +36,16 @@ assert 1 in {1, 2}
 assert 1 in {1: "a"}
 assert 1 in (1, 2)
 assert 1 in map(lambda x: x + 1, [0, 1, 2])
+
+def func(d: dict, t: tuple, s: set):
+    _ = d.get("a")
+    s.add(1)
+    for i in t:
+        print(i)
+
+list_comp = [i + 1 for i in range(10)]
+assert list_comp[0] == 1
+set_comp = {i + 1 for i in range(10)}
+assert 1 in set_comp
+dict_comp = {i: i + 1 for i in range(10)}
+assert dict_comp[0] == 1

tests/test.rs

@@ -165,7 +165,7 @@ fn exec_casting() -> Result<(), String> {
 
 #[test]
 fn exec_collection() -> Result<(), String> {
-    expect("tests/collection.py", 0, 4)
+    expect("tests/collection.py", 0, 5)
 }
 
 #[test]