Implement Statically Typed Dictionaries and Lists

py2v_transpiler/core/analyzer.py

@@ -225,9 +225,59 @@ def analyze(self, tree: ast.AST) -> Dict[str, str]:
 
         return self.type_map
 
+    def _guess_type(self, node: ast.AST) -> str:
+        """Heuristic type guessing for analyzer (simplified version of TranslatorBase._guess_type)."""
+        if isinstance(node, ast.Constant):
+            if isinstance(node.value, bool): return "bool"
+            if isinstance(node.value, int): return "int"
+            if isinstance(node.value, float): return "f64"
+            if isinstance(node.value, str): return "string"
+            if node.value is None: return "int"
+            return "int"
+        elif isinstance(node, (ast.List, ast.Tuple)):
+            if not node.elts: return "[]int"
+            types = [self._guess_type(elt) for elt in node.elts]
+            first = types[0]
+            return f"[]{first}" if all(t == first for t in types) else "[]Any"
+        elif isinstance(node, ast.Dict):
+            if not node.keys: return "map[string]int"
+            k_types = [self._guess_type(k) for k in node.keys if k]
+            v_types = [self._guess_type(v) for v in node.values]
+            k_type = k_types[0] if k_types and all(t == k_types[0] for t in k_types) else "Any"
+            v_type = v_types[0] if v_types and all(t == v_types[0] for t in v_types) else "Any"
+            if k_type == "Any" and k_types and all(isinstance(k, ast.Constant) and isinstance(k.value, str) for k in node.keys if k):
+                k_type = "string"
+            if v_type == "Any" and v_types and all(t != "Any" for t in v_types):
+                # Mixed but not Any? Still Any for now, but if it's mixed classes...
+                pass
+            return f"map[{k_type}]{v_type}"
+        elif isinstance(node, ast.Name):
+            return self.type_map.get(node.id, "Any")
+        elif isinstance(node, ast.Call):
+            if isinstance(node.func, ast.Name):
+                # Common builtins or class names
+                fid = node.func.id
+                if fid == "str": return "string"
+                if fid == "int": return "int"
+                if fid == "float": return "f64"
+                if fid == "bool": return "bool"
+                if fid == "list": return "[]Any"
+                if fid == "dict": return "map[string]Any"
+                return fid
+            elif isinstance(node.func, ast.Attribute):
+                # heuristic: .append() -> returns void, .split() -> []string
+                if node.func.attr == "split": return "[]string"
+                if node.func.attr == "join": return "string"
+                return "Any"
+        return "Any"
+
     def visit_Assign(self, node: ast.Assign) -> Any:
         for target in node.targets:
-            if isinstance(target, ast.Subscript):
+            if isinstance(target, ast.Name):
+                v_type = self._guess_type(node.value)
+                if v_type != "Any":
+                    self.type_map[target.id] = v_type
+            elif isinstance(target, ast.Subscript):
                 dict_name = None
                 if isinstance(target.value, ast.Name):
                     dict_name = target.value.id
@@ -251,32 +301,16 @@ def visit_Assign(self, node: ast.Assign) -> Any:
                         elif isinstance(target.slice.value, str):
                             key_type = "string"
 
-                    val_type = "Any"
-                    if isinstance(node.value, ast.Constant):
-                        if isinstance(node.value.value, int):
-                            val_type = "int"
-                        elif isinstance(node.value.value, str):
-                            val_type = "string"
-                    elif isinstance(node.value, ast.Tuple):
-                        if node.value.elts:
-                            if isinstance(node.value.elts[0], ast.Constant):
-                                if isinstance(node.value.elts[0].value, int):
-                                    val_type = "[]int"
-                                elif isinstance(node.value.elts[0].value, str):
-                                    val_type = "[]string"
-                                else:
-                                    val_type = "[]Any"
-                            else:
-                                val_type = "[]Any"
-                        else:
-                            val_type = "[]Any"
-                    elif isinstance(node.value, ast.Call) and isinstance(
-                        node.value.func, ast.Name
-                    ):
-                        val_type = node.value.func.id
+                    val_type = self._guess_type(node.value)
 
-                    new_type = f"map[{key_type}]{val_type}"
-                    self.type_map[dict_name] = new_type
+                    if dict_name in self.type_map and self.type_map[dict_name].startswith("map["):
+                        # Already known as a map, maybe we can refine it
+                        current_type = self.type_map[dict_name]
+                        if ("Any" in current_type or "int" in current_type) and "Any" not in val_type:
+                             self.type_map[dict_name] = f"map[{key_type}]{val_type}"
+                    else:
+                        new_type = f"map[{key_type}]{val_type}"
+                        self.type_map[dict_name] = new_type
 
         self.generic_visit(node)
 

py2v_transpiler/core/translator/base.py

@@ -393,7 +393,7 @@ def _guess_type(self, node: ast.AST) -> str:
              if isinstance(node.value, float): return "f64"
              if isinstance(node.value, str): return "string"
              if isinstance(node.value, complex): return "PyComplex"
-             if node.value is None: return "int"
+             if node.value is None: return "int" # Fallback to int for uninitialized optionals
              return "int"
         elif isinstance(node, (ast.UnaryOp)):
             if isinstance(node.op, ast.Not):
@@ -412,15 +412,36 @@ def _guess_type(self, node: ast.AST) -> str:
                 if fid == "input": return "string"
                 if fid in ("isinstance", "hasattr", "getattr", "setattr"): return "bool"
         elif isinstance(node, (ast.List, ast.Tuple)):
-            if node.elts:
-                return f"[]{self._guess_type(node.elts[0])}"
-            return "[]int"
+            if not node.elts:
+                return getattr(self, "current_assignment_type", "[]int")
+            types = [self._guess_type(elt) for elt in node.elts]
+            first_type = types[0]
+            if all(t == first_type for t in types):
+                return f"[]{first_type}"
+            return "[]Any"
+        elif isinstance(node, ast.Set):
+            if not node.elts:
+                return "map[int]bool"
+            types = [self._guess_type(elt) for elt in node.elts]
+            first_type = types[0]
+            if all(t == first_type for t in types):
+                return f"map[{first_type}]bool"
+            return "map[Any]bool"
         elif isinstance(node, ast.Dict):
-            if node.keys and node.keys[0]:
-                k_type = self._guess_type(node.keys[0])
-                v_type = self._guess_type(node.values[0])
-                return f"map[{k_type}]{v_type}"
-            return "map[string]int"
+            if not node.keys:
+                return getattr(self, "current_assignment_type", "map[string]int")
+            k_types = [self._guess_type(k) for k in node.keys if k]
+            v_types = [self._guess_type(v) for v in node.values]
+
+            k_type = k_types[0] if k_types else "string"
+            if not all(t == k_type for t in k_types):
+                k_type = "Any"
+
+            v_type = v_types[0] if v_types else "int"
+            if not all(t == v_type for t in v_types):
+                v_type = "Any"
+
+            return f"map[{k_type}]{v_type}"
         elif isinstance(node, ast.Name):
             # Check for location-based type mapping (from mypy plugin)
             if hasattr(node, 'lineno') and hasattr(node, 'col_offset'):

py2v_transpiler/core/translator/classes.py

@@ -683,7 +683,7 @@ def is_enum_type(name: str) -> tuple[bool, bool, bool]:
             # Register that this class has a custom factory
             if not hasattr(self, "defined_classes"):
                 self.defined_classes = {}
-            self.defined_classes[struct_name] = True
+            self.defined_classes[struct_name] = {"has_init": True, "has_new": False}
 
         if is_unittest:
             self.current_class_is_unittest = True
@@ -967,9 +967,9 @@ def is_enum_type(name: str) -> tuple[bool, bool, bool]:
         if not hasattr(self, "defined_classes"):
             self.defined_classes = {}
 
-        # Don't overwrite if it was already set to True (e.g. by dataclass factory)
-        if not self.defined_classes.get(struct_name):
-            self.defined_classes[struct_name] = has_init
+        # Don't overwrite if it was already set (e.g. by dataclass factory)
+        if struct_name not in self.defined_classes:
+            self.defined_classes[struct_name] = {"has_init": has_init, "has_new": False}
 
         # Ensure we output the nested struct definition at the top level
         # visit_ClassDef processes body elements via iteration.

py2v_transpiler/core/translator/literals.py

@@ -73,6 +73,7 @@ def to_standard_str(s):
         return str(val)
 
     def visit_List(self, node: ast.List) -> str:
+        v_type = self._guess_type(node)
         # Check for starred elements
         has_starred = any(isinstance(elt, ast.Starred) for elt in node.elts)
         if has_starred:
@@ -82,19 +83,29 @@ def visit_List(self, node: ast.List) -> str:
             for elt in node.elts:
                 if isinstance(elt, ast.Starred):
                     if current_chunk:
-                        chunks.append(f"[{', '.join(current_chunk)}]")
+                        if v_type == "[]Any":
+                            chunks.append(f"[{', '.join([f'Any({e})' for e in current_chunk])}]")
+                        else:
+                            chunks.append(f"[{', '.join(current_chunk)}]")
                         current_chunk = []
                     chunks.append(str(self.visit(elt.value)))
                 else:
                     current_chunk.append(str(self.visit(elt)))
             if current_chunk:
-                chunks.append(f"[{', '.join(current_chunk)}]")
+                if v_type == "[]Any":
+                    chunks.append(f"[{', '.join([f'Any({e})' for e in current_chunk])}]")
+                else:
+                    chunks.append(f"[{', '.join(current_chunk)}]")
 
             return f"py_list_concat({', '.join(chunks)})"
 
         elements = [str(self.visit(elt)) for elt in node.elts]
         if not elements:
-             return "[]int{}" # Placeholder for empty list
+             return f"{v_type}{{}}"
+
+        if v_type == "[]Any":
+             return f"[{', '.join([f'Any({e})' for e in elements])}]"
+
         return f"[{', '.join(elements)}]"
 
     def visit_Dict(self, node: ast.Dict) -> str:
@@ -124,7 +135,7 @@ def visit_Dict(self, node: ast.Dict) -> str:
                     # Unpacking **expr
                     if current_chunk:
                         # Flush current chunk
-                        chunk_str = f"map[string]int{{{', '.join(current_chunk)}}}"
+                        chunk_str = f"{v_type}{{{', '.join(current_chunk)}}}"
                         chunks.append(chunk_str)
                         current_chunk = []
                     chunks.append(str(self.visit(v)))
@@ -134,56 +145,72 @@ def visit_Dict(self, node: ast.Dict) -> str:
                     current_chunk.append(f"{key_str}: {val_str}")
 
             if current_chunk:
-                chunk_str = f"map[string]int{{{', '.join(current_chunk)}}}"
+                chunk_str = f"{v_type}{{{', '.join(current_chunk)}}}"
                 chunks.append(chunk_str)
 
             return f"py_dict_merge({', '.join(chunks)})"
 
         if not node.keys:
             # Empty dict
-            return "map[string]int{}" # Default fallback
+            return f"{v_type}{{}}"
 
         pairs = []
+        is_any_val = v_type.endswith("Any")
+        is_any_key = "map[Any]" in v_type
+
         for k, v in zip(node.keys, node.values):
             if k:
                 key_str = self.visit(k)
+                if is_any_key:
+                    key_str = f"Any({key_str})"
                 val_str = self.visit(v)
+                if is_any_val:
+                    val_str = f"Any({val_str})"
                 pairs.append(f"{key_str}: {val_str}")
-        return f"map[string]int{{{', '.join(pairs)}}}"
+        return f"{v_type}{{{', '.join(pairs)}}}"
 
     def visit_Set(self, node: ast.Set) -> str:
+        v_type = self._guess_type(node)
         # Check for starred elements
         has_starred = any(isinstance(elt, ast.Starred) for elt in node.elts)
         if has_starred:
             self.used_dict_merge = True
             chunks: List[str] = []
             current_chunk: List[str] = []
+            is_any = "map[Any]" in v_type
             for elt in node.elts:
                 if isinstance(elt, ast.Starred):
                     if current_chunk:
-                        chunks.append(f"map[int]bool{{{', '.join(current_chunk)}}}")
+                        chunk_els = [f"Any({e})" if is_any else e for e in current_chunk]
+                        chunks.append(f"{v_type}{{{', '.join([f'{e}: true' for e in chunk_els])}}}")
                         current_chunk = []
                     chunks.append(str(self.visit(elt.value)))
                 else:
                     val = self.visit(elt)
-                    current_chunk.append(f"{val}: true")
+                    current_chunk.append(val)
 
             if current_chunk:
-                chunks.append(f"map[int]bool{{{', '.join(current_chunk)}}}")
+                chunk_els = [f"Any({e})" if is_any else e for e in current_chunk]
+                chunks.append(f"{v_type}{{{', '.join([f'{e}: true' for e in chunk_els])}}}")
 
             return f"py_dict_merge({', '.join(chunks)})"
 
-        # {1, 2} -> map[int]bool{1: true, 2: true}
-        # Simplified assumption that elements are ints
         elements = []
+        is_any = "map[Any]" in v_type
         for elt in node.elts:
             val = self.visit(elt)
+            if is_any:
+                val = f"Any({val})"
             elements.append(f"{val}: true")
 
-        return f"map[int]bool{{{', '.join(elements)}}}"
+        if not elements:
+            return f"{v_type}{{}}"
+
+        return f"{v_type}{{{', '.join(elements)}}}"
 
     def visit_Tuple(self, node: ast.Tuple) -> str:
         # Translate Tuple (a, b) to Array [a, b]
+        v_type = self._guess_type(node)
         # Check for starred elements
         has_starred = any(isinstance(elt, ast.Starred) for elt in node.elts)
         if has_starred:
@@ -193,17 +220,29 @@ def visit_Tuple(self, node: ast.Tuple) -> str:
             for elt in node.elts:
                 if isinstance(elt, ast.Starred):
                     if current_chunk:
-                        chunks.append(f"[{', '.join(current_chunk)}]")
+                        if v_type == "[]Any":
+                            chunks.append(f"[{', '.join([f'Any({e})' for e in current_chunk])}]")
+                        else:
+                            chunks.append(f"[{', '.join(current_chunk)}]")
                         current_chunk = []
                     chunks.append(str(self.visit(elt.value)))
                 else:
                     current_chunk.append(str(self.visit(elt)))
             if current_chunk:
-                chunks.append(f"[{', '.join(current_chunk)}]")
+                if v_type == "[]Any":
+                    chunks.append(f"[{', '.join([f'Any({e})' for e in current_chunk])}]")
+                else:
+                    chunks.append(f"[{', '.join(current_chunk)}]")
 
             return f"py_list_concat({', '.join(chunks)})"
 
         elements = [str(self.visit(elt)) for elt in node.elts]
+        if not elements:
+            return f"{v_type}{{}}"
+
+        if v_type == "[]Any":
+             return f"[{', '.join([f'Any({e})' for e in elements])}]"
+
         return f"[{', '.join(elements)}]"
 
     def visit_JoinedStr(self, node: ast.JoinedStr) -> str:

py2v_transpiler/tests/translator/test_list_unpacking.py

@@ -18,7 +18,7 @@ def test_list_concat(self):
         source = "x = [1, *a, 2]"
         result = self.transpile(source)
         # Should be py_list_concat([1], a, [2])
-        # Note: visit_List emits [1] as [1], but py_list_concat args.
+        # Note: visit_List emits [1] within py_list_concat
         self.assertIn("py_list_concat", result)
         self.assertIn("([1], a, [2])", result)
 

py2v_transpiler/tests/translator/test_tuples.py

@@ -14,7 +14,9 @@ def test_translator_tuple_literal():
     analyzer.analyze(tree)
     result = translator.visit_Module(tree)
 
-    assert "t := [1, 2]" in result
+    assert "mut t := []int{cap: 2}" in result
+    assert "t << 1" in result
+    assert "t << 2" in result
 
 def test_translator_tuple_mixed():
     parser = PyASTParser()
@@ -29,4 +31,6 @@ def test_translator_tuple_mixed():
     analyzer.analyze(tree)
     result = translator.visit_Module(tree)
 
-    assert "t := [1, 'a']" in result
+    assert "mut t := []Any{cap: 2}" in result
+    assert "t << 1" in result
+    assert "t << 'a'" in result