Types error handling

py2mojo/converters/functiondef.py

@@ -4,6 +4,7 @@
 
 from tokenize_rt import Token
 
+from ..exceptions import ParseException
 from ..helpers import (
     ast_to_offset,
     get_annotation_type,
@@ -63,6 +64,10 @@ def convert_functiondef(node: ast.FunctionDef, rules: RuleSet = 0) -> Iterable:
             )
             continue
 
+        if rules.convert_def_to_fn and not arg.annotation:
+            raise ParseException(
+                node, 'For converting a def function to fn, the declaration needs to be fully type annotated'
+            )
         curr_type = get_annotation_type(arg.annotation)
         new_type = get_mojo_type(curr_type)
 

py2mojo/exceptions.py

@@ -0,0 +1,7 @@
+import ast
+
+
+class ParseException(Exception):
+    def __init__(self, node: ast.AST, msg: str):
+        self.node = node
+        self.msg = msg

py2mojo/helpers.py

@@ -1,6 +1,9 @@
 import ast
 import re
 
+import astor
+from rich import print
+from rich.text import Text
 from tokenize_rt import UNIMPORTANT_WS, Offset, Token
 
 
@@ -115,3 +118,29 @@ def get_mojo_type(curr_type: str) -> str:
             curr_type = pattern.sub(replacement, curr_type)
 
     return curr_type
+
+
+def highlight_code_at_position(code: str, line: int, column: int, end_column: int) -> Text:
+    lines = code.splitlines()
+    highlighted = Text()
+
+    for idx, source_line in enumerate(lines):
+        if idx + 1 == line:
+            # Highlight the specific column in the given line
+            highlighted.append(source_line[:column], style='white')
+            for i in range(column, min(end_column, len(source_line))):
+                highlighted.append(source_line[i], style='bold black on yellow')
+            highlighted.append(source_line[end_column + 1 :], style='white')
+        else:
+            highlighted.append(source_line, style='white')
+        highlighted.append('\n')
+
+    return highlighted
+
+
+def display_error(node: ast.AST, message: str):
+    src = astor.to_source(node)
+
+    highlighted_src = highlight_code_at_position(src, 1, node.col_offset, node.end_col_offset)
+    print('[bold red]Error:[/bold red]', message)
+    print(highlighted_src)

py2mojo/main.py

@@ -11,7 +11,8 @@
 from tokenize_rt import Token, reversed_enumerate, src_to_tokens, tokens_to_src
 
 from .converters import convert_assignment, convert_functiondef, convert_classdef
-from .helpers import fixup_dedent_tokens
+from .exceptions import ParseException
+from .helpers import display_error, fixup_dedent_tokens
 from .rules import get_rules, RuleSet
 
 
@@ -113,7 +114,11 @@ def main(argv: Sequence[str] | None = None) -> int:
 
             rules = get_rules(args)
 
-            annotated_source = convert_to_mojo(source, rules)
+            try:
+                annotated_source = convert_to_mojo(source, rules)
+            except ParseException as exc:
+                display_error(exc.node, exc.msg)
+                sys.exit(1)
 
             if source != annotated_source:
                 print(f'Rewriting {filename}' if args.inplace else f'Rewriting {filename} into {mojo_filename}')

requirements.txt

@@ -1 +1,3 @@
-tokenize_rt>=5.2.0
+tokenize_rt>=5.2.0
+rich>=13.5.2
+astor>=0.8.1

tests/conftest.py

@@ -0,0 +1,4 @@
+import pytest
+
+# We want pytest assert introspection in the helpers
+pytest.register_assert_rewrite('helpers')

tests/test_functiondef.py

@@ -1,6 +1,7 @@
 import pytest
 
 from helpers import validate
+from py2mojo.exceptions import ParseException
 from py2mojo.rules import RuleSet
 
 
@@ -72,3 +73,16 @@ class Point:
     def __init__(inout self, x: Int, y: Int) -> Int: ...
 ''',
     )
+
+
+def test_functiondef_non_fully_annotated_functions():
+    validate(
+        '''def add(x, y): return x + y''',
+        '''def add(x, y): return x + y''',
+    )
+    with pytest.raises(ParseException):
+        validate(
+            '''def add(x, y): return x + y''',
+            '''def add(x, y): return x + y''',
+            rules=RuleSet(convert_def_to_fn=True),
+        )