Convert the str, dict and list types properly to equivalent in Mojo

msaelices · May 18, 2024 · d1a04f0 · d1a04f0
1 parent d6d9b2f
commit d1a04f0
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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
     )