From 7882a874fb34a64110673ead263785e8c3db87d0 Mon Sep 17 00:00:00 2001 From: Manuel Saelices Date: Thu, 14 Sep 2023 23:37:53 +0200 Subject: [PATCH 1/2] We want pytest assert introspection in the helpers --- tests/conftest.py | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 tests/conftest.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..7f92049 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,4 @@ +import pytest + +# We want pytest assert introspection in the helpers +pytest.register_assert_rewrite('helpers') From 3b6eace4c3c1ce5afcadda3fb1b4ea1c20fb619d Mon Sep 17 00:00:00 2001 From: Manuel Saelices Date: Fri, 15 Sep 2023 00:36:49 +0200 Subject: [PATCH 2/2] Detect non-fully-annotated "def" Python functions when we try to convert to "fn", logging the error --- py2mojo/converters/functiondef.py | 5 +++++ py2mojo/exceptions.py | 7 +++++++ py2mojo/helpers.py | 29 +++++++++++++++++++++++++++++ py2mojo/main.py | 9 +++++++-- requirements.txt | 4 +++- tests/test_functiondef.py | 14 ++++++++++++++ 6 files changed, 65 insertions(+), 3 deletions(-) create mode 100644 py2mojo/exceptions.py diff --git a/py2mojo/converters/functiondef.py b/py2mojo/converters/functiondef.py index f0b1e1b..c5470c8 100644 --- a/py2mojo/converters/functiondef.py +++ b/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) diff --git a/py2mojo/exceptions.py b/py2mojo/exceptions.py new file mode 100644 index 0000000..d95f7ac --- /dev/null +++ b/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 diff --git a/py2mojo/helpers.py b/py2mojo/helpers.py index a8dc44a..d0cf53a 100644 --- a/py2mojo/helpers.py +++ b/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) diff --git a/py2mojo/main.py b/py2mojo/main.py index 0fe6e6e..5d04c0b 100644 --- a/py2mojo/main.py +++ b/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}') diff --git a/requirements.txt b/requirements.txt index 5e29958..0950931 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,3 @@ -tokenize_rt>=5.2.0 \ No newline at end of file +tokenize_rt>=5.2.0 +rich>=13.5.2 +astor>=0.8.1 \ No newline at end of file diff --git a/tests/test_functiondef.py b/tests/test_functiondef.py index 6c057c1..7f1f049 100644 --- a/tests/test_functiondef.py +++ b/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), + )