phasm/py2wasm/python.py

"""
Code for parsing the source (python-alike) code
"""

from typing import List, Optional, Union, Tuple

import ast

from . import wasm

# pylint: disable=C0103,R0201

class Visitor:
    """
    Class to visit a Python syntax tree

    Since we need to visit the whole tree, there's no point in subclassing
    the buildin visitor
    """
    def visit_Module(self, node: ast.Module) -> wasm.Module:
        """
        Visits a Python module, which results in a wasm Module
        """
        assert not node.type_ignores

        module = wasm.Module()

        for stmt in node.body:
            if isinstance(stmt, ast.FunctionDef):
                wnode = self.visit_FunctionDef(stmt)

                if isinstance(wnode, wasm.Import):
                    module.imports.append(wnode)
                else:
                    module.functions.append(wnode)
                continue

            raise NotImplementedError(stmt)

        return module

    def visit_FunctionDef(self, node: ast.FunctionDef) -> Union[wasm.Import, wasm.Function]:
        """
        A Python function definition with the @external decorator is returned as import.

        Other functions are returned as (exported) functions.
        """

        exported = False

        if node.decorator_list:
            assert 1 == len(node.decorator_list)

            decorator = node.decorator_list[0]
            if isinstance(decorator, ast.Name):
                assert 'exported' == decorator.id
                exported = True
            elif isinstance(decorator, ast.Call):
                call = decorator

                assert not node.type_comment
                assert 1 == len(node.body)
                assert isinstance(node.body[0], ast.Expr)
                assert isinstance(node.body[0].value, ast.Ellipsis)

                assert isinstance(call.func, ast.Name)
                assert 'imported' == call.func.id

                assert not call.keywords

                module, name, intname = _parse_import_decorator(node.name, call.args)

                return wasm.Import(module, name, intname, _parse_import_args(node.args))

        return wasm.Function(node.name, exported)

def _parse_import_decorator(func_name: str, args: List[ast.expr]) -> Tuple[str, str, str]:
    """
    Parses an @import decorator
    """

    assert 0 < len(args) < 3
    str_args = [
        arg.value
        for arg in args
        if isinstance(arg, ast.Constant)
        and isinstance(arg.value, str)
    ]
    assert len(str_args) == len(args)

    module = str_args.pop(0)
    if str_args:
        name = str_args.pop(0)
    else:
        name = func_name

    return module, name, func_name

def _parse_import_args(args: ast.arguments) -> List[Tuple[str, str]]:
    """
    Parses the arguments for an @imported method
    """

    assert not args.vararg
    assert not args.kwonlyargs
    assert not args.kw_defaults
    assert not args.kwarg
    assert not args.defaults # Maybe support this later on

    arg_list = [
        *args.posonlyargs,
        *args.args,
    ]

    return [
        (arg.arg, _parse_annotation(arg.annotation))
        for arg in arg_list
    ]

def _parse_annotation(ann: Optional[ast.expr]) -> str:
    """
    Parses a typing annotation
    """
    assert ann is not None, 'Web Assembly requires type annotations'
    assert isinstance(ann, ast.Name)
    result = ann.id
    assert result in ['i32']
    return result




    # def visit_ImportFrom(self, node):
    #     for alias in node.names:
    #         self.imports.append(Import(
    #             node.module,
    #             alias.name,
    #             alias.asname,
    #         ))
    #
    # def generic_visit(self, node: Any) -> None:
    #     print(node)
    #     super().generic_visit(node)
    #
    # def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
    #     is_export = False
    #
    #     if node.decorator_list:
    #         assert 1 == len(node.decorator_list)
    #
    #         call = node.decorator_list[0]
    #         if not isinstance(call, ast.Name):
    #             assert isinstance(call, ast.Call)
    #
    #             assert 'external' == call.func.id
    #             assert 1 == len(call.args)
    #             assert isinstance(call.args[0].value, str)
    #
    #             import_ = Import(
    #                 call.args[0].value,
    #                 node.name,
    #                 node.name,
    #             )
    #
    #             import_.params = [
    #                 arg.annotation.id
    #                 for arg in node.args.args
    #             ]
    #
    #             self.module.imports.append(import_)
    #             return
    #
    #         assert call.id == 'export'
    #         is_export = True
    #
    #     func = Function(
    #         node.name,
    #         is_export,
    #     )
    #
    #     for arg in node.args.args:
    #         func.params.append(
    #             (arg.arg, arg.annotation.id)
    #         )
    #     func.result = node.returns
    #
    #     self._stack.append(func)
    #     self.generic_visit(node)
    #     self._stack.pop()
    #
    #     self.module.functions.append(func)
    #
    # def visit_Call(self, node: ast.Call) -> None:
    #     self.generic_visit(node)
    #
    #     func = self._stack[-1]
    #     func.statements.append(
    #         Statement('call', '$' + node.func.id)
    #     )
    #
    # def visit_BinOp(self, node: ast.BinOp) -> None:
    #     self.generic_visit(node)
    #
    #     func = self._stack[-1]
    #
    #     if 'Add' == node.op.__class__.__name__:
    #         func.statements.append(
    #             Statement('i32.add')
    #         )
    #     elif 'Mult' == node.op.__class__.__name__:
    #         func.statements.append(
    #             Statement('i32.mul')
    #         )
    #     else:
    #         raise NotImplementedError
    #
    # def visit_Constant(self, node: ast.Constant) -> None:
    #     func = self._stack[-1]
    #     if isinstance(node.value, int):
    #         func.statements.append(
    #             Statement('i32.const', str(node.value))
    #         )
    #
    #     self.generic_visit(node)