Automated Python to Mojo code translation
from .assignment import convert_assignment
from .classdef import convert_classdef
from .functiondef import convert_functiondef

__all__ = [
    'convert_assignment',
    'convert_classdef',
    'convert_functiondef',
]

py2mojo/converters/__init__.py
msaelices/py2mojo
import ast
from functools import partial
from typing import Iterable

from tokenize_rt import Token

from ..helpers import ast_to_offset, get_annotation_type, find_token, find_token_by_name, get_mojo_type
from ..rules import RuleSet


def _replace_assignment(tokens: list[Token], i: int, rules: RuleSet, new_type: str) -> None:
    tokens.insert(i, Token(name='NAME', src='var '))
    ann_idx = find_token(tokens, i, ':')
    type_idx = find_token_by_name(tokens, ann_idx, name='NAME')
    assign_op_idx = find_token(tokens, type_idx, '=') - 1
    newline_idx = find_token_by_name(tokens, type_idx, name='NEWLINE')
    valid_idxs = (idx for idx in (assign_op_idx, newline_idx) if idx >= 0)
    end_type_idx = min(valid_idxs)
    del tokens[type_idx:end_type_idx]
    tokens.insert(type_idx, Token(name='NAME', src=new_type))


def convert_assignment(node: ast.AnnAssign, rules: RuleSet) -> Iterable:
    """Convert an assignment to a mojo assignment."""
    curr_type = get_annotation_type(node.annotation)
    new_type = get_mojo_type(curr_type, rules)
    if not new_type:
        return

    yield (
        ast_to_offset(node),
        partial(
            _replace_assignment,
            new_type=new_type,
        ),
    )

py2mojo/converters/assignment.py
msaelices/py2mojo
import ast
from functools import partial
from typing import Iterable

from tokenize_rt import Token

from ..exceptions import ParseException
from ..helpers import ast_to_offset, find_token
from ..rules import RuleSet


def _replace_class_keyword(tokens: list, i: int, rules: RuleSet) -> None:
    idx = find_token(tokens, i, 'class')
    tokens[idx] = Token(name='NAME', src='struct')


def convert_classdef(node: ast.FunctionDef, rules: RuleSet) -> Iterable:
    """Converts the annotation of the given function definition."""
    if rules.convert_class_to_struct:
        for child in node.body:
            if isinstance(child, ast.Assign):
                raise ParseException(node, 'Class contains non type annotated attributes.')

        offset = ast_to_offset(node)
        yield (
            offset,
            partial(
                _replace_class_keyword,
            ),
        )

py2mojo/converters/classdef.py
msaelices/py2mojo
import ast
from functools import partial
from typing import Iterable

from tokenize_rt import Token

from ..exceptions import ParseException
from ..helpers import (
    ast_to_offset,
    get_annotation_type,
    find_token,
    find_token_after_offset,
    find_token_by_name,
    get_mojo_type,
)
from ..rules import RuleSet


def _replace_annotation(
    tokens: list, i: int, rules: RuleSet, end_offset: int, new_type: str, ann_offset: int | None = None
) -> None:
    if ann_offset:
        ann_idx = find_token_after_offset(tokens, i, ann_offset)
    else:
        ann_idx = find_token(tokens, i, ':')
    type_idx = find_token_by_name(tokens, ann_idx, name='NAME')
    end_type_idx = find_token_after_offset(tokens, ann_idx, end_offset)
    del tokens[type_idx:end_type_idx]
    tokens.insert(type_idx, Token(name='NAME', src=new_type))


def _replace_def_keyword(tokens: list, i: int, rules: RuleSet) -> None:
    idx = find_token(tokens, i, 'def')
    tokens[idx] = Token(name='NAME', src='fn')


def _add_declaration(tokens: list, i: int, rules: RuleSet, declaration: str) -> None:
    tokens.insert(i, Token(name='NAME', src=declaration))
    tokens.insert(i + 1, Token(name='UNIMPORTANT_WS', src=' '))


def convert_functiondef(node: ast.FunctionDef, rules: RuleSet = 0) -> Iterable:
    """Converts the annotation of the given function definition."""
    if rules.convert_def_to_fn > 0:
        offset = ast_to_offset(node)
        yield (
            offset,
            partial(
                _replace_def_keyword,
            ),
        )

    if not node.args.args:
        return

    for arg in node.args.args:
        if arg.arg == 'self':
            yield (
                ast_to_offset(arg),
                partial(
                    _add_declaration,
                    declaration='inout',
                ),
            )
            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, rules)

        if not new_type:
            continue

        yield (
            ast_to_offset(arg),
            partial(
                _replace_annotation,
                end_offset=arg.end_col_offset,
                new_type=new_type,
            ),
        )

    if node.returns:
        curr_type = get_annotation_type(node.returns)
        new_type = get_mojo_type(curr_type, rules)
        if not new_type:
            return

        offset = ast_to_offset(node.returns)
        yield (
            offset,
            partial(
                _replace_annotation,
                end_offset=node.returns.end_col_offset,
                new_type=new_type,
                ann_offset=offset.utf8_byte_offset,
            ),
        )

py2mojo/converters/functiondef.py
msaelices/py2mojo
import ast


class ParseException(Exception):
    def __init__(self, node: ast.AST, msg: str):
        self.node = node
        self.msg = msg

py2mojo/exceptions.py
msaelices/py2mojo
import ast
import re

import astor
from rich import print
from rich.text import Text
from tokenize_rt import UNIMPORTANT_WS, Offset, Token

from .rules import RuleSet


def ast_to_offset(node: ast.expr | ast.stmt) -> Offset:
    return Offset(node.lineno, node.col_offset)


def find_token(tokens: list[Token], i: int, src: str) -> int:
    """Find the index of the token with the given src."""
    try:
        while tokens[i].src != src:
            i += 1
    except IndexError:
        return -1
    return i


def find_token_by_name(tokens: list[Token], i: int, name: str) -> int:
    """Find the index of the token with the given name."""
    try:
        while tokens[i].name != name:
            i += 1
    except IndexError:
        return -1
    return i


def find_token_after_offset(tokens: list[Token], i: int, offset: int) -> int:
    """Find the index of the token after the given offset."""
    try:
        while tokens[i].utf8_byte_offset < offset:
            i += 1
    except IndexError:
        return -1
    return i


def fixup_dedent_tokens(tokens: list[Token]) -> None:
    # copied from pyupgrade
    """For whatever reason the DEDENT / UNIMPORTANT_WS tokens are misordered
    | if True:
    |     if True:
    |         pass
    |     else:
    |^    ^- DEDENT
    |+----UNIMPORTANT_WS
    """
    for i, token in enumerate(tokens):
        if token.name == UNIMPORTANT_WS and tokens[i + 1].name == 'DEDENT':
            tokens[i], tokens[i + 1] = tokens[i + 1], tokens[i]


def get_node_name(node: ast.AST) -> str:
    """Returns the name of the given node."""
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Constant):
        return node.value
    elif isinstance(node, ast.Attribute):
        return node.attr
    elif isinstance(node, ast.Call):
        return node.func.id if isinstance(node.func, ast.Name) else node.func.attr
    else:
        return 'unknown'


def get_next_node(node: ast.AST) -> ast.AST:
    """Returns the next node of the given node."""
    if isinstance(node, ast.Attribute):
        return node.value
    elif isinstance(node, ast.Call):
        return node.func if isinstance(node.func, ast.Name) else node.func.value
    else:
        raise ValueError(f'Unexpected node type: {type(node)}')


def get_dot_path(node: ast.Attribute):
    """Returns the dot path of the given attribute node."""
    attr_list = []
    while isinstance(node, ast.Attribute) or isinstance(node, ast.Call):
        attr = get_node_name(node)
        node = get_next_node(node)
        attr_list.append(attr)
    attr_list.append(get_node_name(node))
    return reversed(attr_list)


def get_annotation_type(node: ast.AST) -> str:
    """Returns the type of the given annotation node."""
    match node.__class__.__name__:
        case 'Name':
            curr_type = node.id
        case 'Subscript':
            curr_type = f'{node.value.id}[{get_annotation_type(node.slice)}]'
        case _:
            curr_type = ''
    return curr_type


def get_mojo_type(curr_type: str, rules: RuleSet) -> str:
    """Returns the corresponding Mojo type for the given Python type."""
    patterns = [
        (re.compile(r'int'), 'Int'),
        (re.compile(r'float'), f'Float{rules.float_precision}'),
    ]

    prev_type = ''

    while prev_type != curr_type:
        prev_type = curr_type
        for pattern, replacement in patterns:
            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('\n[bold red]Error:[/bold red]', message)
    print('[light_salmon3]Source[/light_salmon3]:\n', highlighted_src)

py2mojo/helpers.py
msaelices/py2mojo
from __future__ import annotations

import argparse
import ast
import os
import sys
import tokenize
from collections import defaultdict
from typing import Callable, Sequence

from tokenize_rt import Token, reversed_enumerate, src_to_tokens, tokens_to_src

from .converters import convert_assignment, convert_functiondef, convert_classdef
from .exceptions import ParseException
from .helpers import display_error, fixup_dedent_tokens
from .rules import get_rules, RuleSet


TokenFunc = Callable[[list[Token], int], None]


def get_converters(klass: type) -> list[TokenFunc]:
    return {
        ast.ClassDef: [
            convert_classdef,
        ],
        ast.AnnAssign: [
            convert_assignment,
        ],
        ast.FunctionDef: [
            convert_functiondef,
        ],
    }.get(klass, [])


def visit(tree: ast.Module, rules: RuleSet) -> list[TokenFunc]:
    nodes = [tree]
    ret = defaultdict(list)
    while nodes:
        node = nodes.pop()

        for converter in get_converters(type(node)):
            for offset, token_func in converter(node, rules):
                ret[offset].append(token_func)

        for name in reversed(node._fields):
            value = getattr(node, name)

            if isinstance(value, ast.AST):
                nodes.append(value)
            elif isinstance(value, list):
                for subvalue in reversed(value):
                    if isinstance(subvalue, ast.AST):
                        nodes.append(subvalue)
    return ret


def convert_to_mojo(source: str, rules: RuleSet) -> str:
    tree = ast.parse(source)

    callbacks = visit(tree, rules)

    if not callbacks:
        return source

    try:
        tokens = src_to_tokens(source)
    except tokenize.TokenError:
        return source

    fixup_dedent_tokens(tokens)

    for i, token in reversed_enumerate(tokens):
        if not token.src:
            continue

        for callback in callbacks.get(token.offset, ()):
            callback(tokens, i, rules)

    return tokens_to_src(tokens)


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument('filenames', nargs='+')
    parser.add_argument(
        '--inplace',
        help='Rewrite the file inplace',
        action='store_true',
    )
    parser.add_argument(
        '--extension',
        help='File extension of the generated files',
        choices=['mojo', '🔥'],
        default='🔥',
        type=str,
    )
    parser.add_argument(
        '--convert-def-to-fn',
        default=True,
        action=argparse.BooleanOptionalAction,
    )
    parser.add_argument(
        '--convert-class-to-struct',
        default=True,
        action=argparse.BooleanOptionalAction,
    )
    parser.add_argument(
        '--float-precision',
        default=32,
        type=int,
        choices=[32, 64],
    )
    args = parser.parse_args(argv)

    for filename in args.filenames:
        mojo_filename = filename if args.inplace else f'{os.path.splitext(filename)[0]}.{args.extension}'
        with open(filename) as source_file:
            source = source_file.read()

            rules = get_rules(args)

            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}')
                with open(mojo_filename, 'w', encoding='UTF-8', newline='') as out:
                    out.write(annotated_source)
            else:
                print(f'File {filename} unchanged')

    return 0


if __name__ == '__main__':
    sys.exit(main())

py2mojo/main.py
msaelices/py2mojo
import argparse
from dataclasses import dataclass


@dataclass
class RuleSet:
    convert_def_to_fn: bool = False
    convert_class_to_struct: bool = False
    float_precision: int = 64


def get_rules(args: argparse.Namespace) -> RuleSet:
    return RuleSet(
        convert_def_to_fn=args.convert_def_to_fn,
        convert_class_to_struct=args.convert_class_to_struct,
        float_precision=args.float_precision,
    )

py2mojo/rules.py
msaelices/py2mojo
import pytest

# We want pytest assert introspection in the helpers
pytest.register_assert_rewrite('helpers')

tests/conftest.py
msaelices/py2mojo
from py2mojo.main import convert_to_mojo
from py2mojo.rules import RuleSet


def validate(source: str, expected: str, rules: RuleSet | None = None) -> None:
    rules = rules or RuleSet()
    converted_source = convert_to_mojo(source, rules=rules)
    assert converted_source == expected

tests/helpers.py
msaelices/py2mojo
from helpers import validate
from py2mojo.rules import RuleSet


def test_fibonacci():
    validate(
        '''
def fib(n: int) -> int:
    """
    Calculate the nth Fibonacci number.
    """
    if n <= 1:
        return n
    else:
        return fib(n - 1) + fib(n - 2)
''',
        '''
def fib(n: Int) -> Int:
    """
    Calculate the nth Fibonacci number.
    """
    if n <= 1:
        return n
    else:
        return fib(n - 1) + fib(n - 2)
''',
    )
    validate(
        '''
def fib(n: int) -> int:
    """
    Calculate the nth Fibonacci number.
    """
    if n <= 1:
        return n
    else:
        return fib(n - 1) + fib(n - 2)
''',
        '''
fn fib(n: Int) -> Int:
    """
    Calculate the nth Fibonacci number.
    """
    if n <= 1:
        return n
    else:
        return fib(n - 1) + fib(n - 2)
''',
        rules=RuleSet(convert_def_to_fn=True),
    )

tests/test_algorithms.py
msaelices/py2mojo
from helpers import validate
from py2mojo.rules import RuleSet


def test_assignment_with_basic_types():
    validate(
        'x: int = 10',
        'var x: Int = 10',
    )
    validate(
        'x: float = 10.5',
        'var x: Float64 = 10.5',
    )
    validate(
        'x: float = 10.5',
        'var x: Float32 = 10.5',
        rules=RuleSet(float_precision=32),
    )
    validate(
        'x: str = "foo"',
        'var x: str = "foo"',
    )
    validate(
        'x: int',
        'var x: Int',
    )
    validate(
        '"""docstring"""\nx: int',
        '"""docstring"""\nvar x: Int',
    )


def test_assignment_with_list_types():
    validate(
        'x: List[int] = []',
        'var x: List[Int] = []',
    )
    validate(
        'x: list[float] = []',
        'var x: list[Float64] = []',
    )
    validate(
        'x: list = []',
        'var x: list = []',
    )

tests/test_assignment.py
msaelices/py2mojo
import pytest

from helpers import validate
from py2mojo.exceptions import ParseException
from py2mojo.rules import RuleSet


def test_classdef():
    validate(
        'class Foo(Bar): pass',
        'class Foo(Bar): pass',
    )
    validate(
        'class Foo(Bar): pass',
        'struct Foo(Bar): pass',
        rules=RuleSet(convert_class_to_struct=True),
    )
    validate(
        'class Foo(Bar): n: int = 10',
        'struct Foo(Bar): var n: Int = 10',
        rules=RuleSet(convert_class_to_struct=True),
    )


def test_classdef_non_fully_annotated_classes():
    validate(
        '''class Number: number = 10''',
        '''class Number: number = 10''',
        rules=RuleSet(convert_class_to_struct=False),
    )
    with pytest.raises(ParseException):
        validate(
            '''class Number: number = 10''',
            '''class Number: number = 10''',
            rules=RuleSet(convert_class_to_struct=True),
        )

tests/test_classdef.py
msaelices/py2mojo
import pytest

from helpers import validate
from py2mojo.exceptions import ParseException
from py2mojo.rules import RuleSet


def test_functiondef_with_no_params():
    validate(
        'def main(): print("Hello world!")',
        'def main(): print("Hello world!")',
    )
    validate(
        'def main(): print("Hello world!")',
        'fn main(): print("Hello world!")',
        rules=RuleSet(convert_def_to_fn=True),
    )


@pytest.mark.parametrize(
    'python_type, mojo_type',
    [
        ('int', 'Int'),
        ('float', 'Float64'),
    ],
)
def test_functiondef_with_basic_types(python_type, mojo_type):
    validate(
        f'def add(x: {python_type}, y: {python_type}) -> {python_type}: return x + y',
        f'def add(x: {mojo_type}, y: {mojo_type}) -> {mojo_type}: return x + y',
    )
    validate(
        f'def add(x: {python_type}, y: {python_type}) -> {python_type}: return x + y',
        f'fn add(x: {mojo_type}, y: {mojo_type}) -> {mojo_type}: return x + y',
        rules=RuleSet(convert_def_to_fn=True),
    )


@pytest.mark.parametrize(
    'python_type, mojo_type',
    [
        ('int', 'Int'),
        ('float', 'Float64'),
    ],
)
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}]: ...',
    )
    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)',
    )
    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',
    )
    validate(
        'def concat(l1: list, l2: list) -> list: return l1 + l2',
        'def concat(l1: list, l2: list) -> list: return l1 + l2',  # no changed
    )


def test_functiondef_with_float_in_precision():
    validate(
        'def add(x: float, y: float) -> float: return x + y',
        'fn add(x: Float32, y: Float32) -> Float32: return x + y',
        rules=RuleSet(convert_def_to_fn=True, float_precision=32),
    )


def test_functiondef_inside_classes():
    validate(
        '''
class Point:
    def __init__(self, x: int, y: int) -> int: ...
''',
        '''
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),
        )

tests/test_functiondef.py
msaelices/py2mojo