From d1a04f06e510e6a435cb313fad07faec254bb2b1 Mon Sep 17 00:00:00 2001
From: Manuel Saelices <msaelices@gmail.com>
Date: Sat, 18 May 2024 15:36:14 +0200
Subject: [PATCH] Convert the str, dict and list types properly to equivalent
 in Mojo

---
 py2mojo/helpers.py        |  3 +++
 tests/test_assignment.py  | 12 ++++++++----
 tests/test_functiondef.py |  8 ++++----
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/py2mojo/helpers.py b/py2mojo/helpers.py
index 76e0545..84f9b19 100644
--- a/py2mojo/helpers.py
+++ b/py2mojo/helpers.py
@@ -110,6 +110,9 @@ def get_mojo_type(curr_type: str, rules: RuleSet) -> str:
     patterns = [
         (re.compile(r'int'), 'Int'),
         (re.compile(r'float'), f'Float{rules.float_precision}'),
+        (re.compile(r'str'), 'String'),
+        (re.compile(r'list'), 'List'),
+        (re.compile(r'dict'), 'Dict'),
     ]
 
     prev_type = ''
diff --git a/tests/test_assignment.py b/tests/test_assignment.py
index 5db81c1..51f5623 100644
--- a/tests/test_assignment.py
+++ b/tests/test_assignment.py
@@ -18,7 +18,7 @@ def test_assignment_with_basic_types():
     )
     validate(
         'x: str = "foo"',
-        'var x: str = "foo"',
+        'var x: String = "foo"',
     )
     validate(
         'x: int',
@@ -28,18 +28,22 @@ def test_assignment_with_basic_types():
         '"""docstring"""\nx: int',
         '"""docstring"""\nvar x: Int',
     )
+    validate(
+        'd: dict = {}',
+        'var d: Dict = {}',
+    )
 
 
 def test_assignment_with_list_types():
     validate(
-        'x: List[int] = []',
+        'x: list[int] = []',
         'var x: List[Int] = []',
     )
     validate(
         'x: list[float] = []',
-        'var x: list[Float64] = []',
+        'var x: List[Float64] = []',
     )
     validate(
         'x: list = []',
-        'var x: list = []',
+        'var x: List = []',
     )
diff --git a/tests/test_functiondef.py b/tests/test_functiondef.py
index 15ce397..29a1924 100644
--- a/tests/test_functiondef.py
+++ b/tests/test_functiondef.py
@@ -46,19 +46,19 @@ def test_functiondef_with_basic_types(python_type, mojo_type):
 def test_functiondef_with_list_types(python_type, mojo_type):
     validate(
         f'def flatten(l1: list[list[{python_type}]]) -> list[{python_type}]: ...',
-        f'def flatten(l1: list[list[{mojo_type}]]) -> list[{mojo_type}]: ...',
+        f'def flatten(l1: List[List[{mojo_type}]]) -> List[{mojo_type}]: ...',
     )
     validate(
         f'def reverse(l: list[{python_type}]) -> list[{python_type}]: return reversed(l)',
-        f'def reverse(l: list[{mojo_type}]) -> list[{mojo_type}]: return reversed(l)',
+        f'def reverse(l: List[{mojo_type}]) -> List[{mojo_type}]: return reversed(l)',
     )
     validate(
         f'def concat(l1: list[{python_type}], l2: list[{python_type}]) -> {python_type}: return l1 + l2',
-        f'def concat(l1: list[{mojo_type}], l2: list[{mojo_type}]) -> {mojo_type}: return l1 + l2',
+        f'def concat(l1: List[{mojo_type}], l2: List[{mojo_type}]) -> {mojo_type}: return l1 + l2',
     )
     validate(
         'def concat(l1: list, l2: list) -> list: return l1 + l2',
-        'def concat(l1: list, l2: list) -> list: return l1 + l2',  # no changed
+        'def concat(l1: List, l2: List) -> List: return l1 + l2',  # no changed
     )