src/basilisp/lang/compiler/generator.py
def __multi_arity_fn_to_py_ast(  # pylint: disable=too-many-locals
    ctx: GeneratorContext,
    node: Fn,
    methods: Collection[FnMethod],
    def_name: Optional[str] = None,
    meta_node: Optional[MetaNode] = None,
) -> GeneratedPyAST:
    """Return a Python AST node for a function with multiple arities."""
    assert node.op == NodeOp.FN
    assert all([method.op == NodeOp.FN_METHOD for method in methods])

    lisp_fn_name = node.local.name if node.local is not None else None
    py_fn_name = __fn_name(lisp_fn_name) if def_name is None else munge(def_name)

    py_fn_node = ast.AsyncFunctionDef if node.is_async else ast.FunctionDef

    arity_to_name = {}
    rest_arity_name: Optional[str] = None
    fn_defs = []
    for method in methods:
        arity_name = f"{py_fn_name}__arity{'_rest' if method.is_variadic else method.fixed_arity}"
        if method.is_variadic:
            rest_arity_name = arity_name
        else:
            arity_to_name[method.fixed_arity] = arity_name

        with ctx.new_symbol_table(arity_name), ctx.new_recur_point(
            method.loop_id, RecurType.FN, is_variadic=node.is_variadic
        ):
            # Allow named anonymous functions to recursively call themselves
            if lisp_fn_name is not None:
                ctx.symbol_table.new_symbol(
                    sym.symbol(lisp_fn_name), py_fn_name, LocalType.FN
                )

            fn_args, varg, fn_body_ast = __fn_args_to_py_ast(
                ctx, method.params, method.body
            )
            fn_defs.append(
                py_fn_node(
                    name=arity_name,
                    args=ast.arguments(
                        args=fn_args,
                        kwarg=None,
                        vararg=varg,
                        kwonlyargs=[],
                        defaults=[],
                        kw_defaults=[],
                    ),
                    body=fn_body_ast,
                    decorator_list=[_TRAMPOLINE_FN_NAME]
                    if ctx.recur_point.has_recur
                    else [],
                    returns=None,
                )
            )

    dispatch_fn_ast = __multi_arity_dispatch_fn(
        ctx,
        py_fn_name,
        arity_to_name,
        default_name=rest_arity_name,
        max_fixed_arity=node.max_fixed_arity,
        meta_node=meta_node,
        is_async=node.is_async,
    )

    return GeneratedPyAST(
        node=dispatch_fn_ast.node,
        dependencies=list(chain(fn_defs, dispatch_fn_ast.dependencies)),
    )
def __multi_arity_fn_to_py_ast(  # pylint: disable=too-many-locals
    ctx: GeneratorContext,
    node: Fn,
    methods: Collection[FnMethod],
    def_name: Optional[str] = None,
    meta_node: Optional[MetaNode] = None,
) -> GeneratedPyAST:
    assert node.op == NodeOp.FN
    assert all([method.op == NodeOp.FN_METHOD for method in methods])

    lisp_fn_name = node.local.name if node.local is not None else None
    py_fn_name = __fn_name(lisp_fn_name) if def_name is None else munge(def_name)

    py_fn_node = ast.AsyncFunctionDef if node.is_async else ast.FunctionDef

    arity_to_name = {}
    rest_arity_name: Optional[str] = None
    fn_defs = []
    for method in methods:
        arity_name = f"{py_fn_name}__arity{'_rest' if method.is_variadic else method.fixed_arity}"
        if method.is_variadic:
            rest_arity_name = arity_name
        else:
            arity_to_name[method.fixed_arity] = arity_name

        with ctx.new_symbol_table(arity_name), ctx.new_recur_point(
            method.loop_id, RecurType.FN, is_variadic=node.is_variadic
        ):
            # Allow named anonymous functions to recursively call themselves
            if lisp_fn_name is not None:
                ctx.symbol_table.new_symbol(
                    sym.symbol(lisp_fn_name), py_fn_name, LocalType.FN
                )

            fn_args, varg, fn_body_ast = __fn_args_to_py_ast(
                ctx, method.params, method.body
            )
            fn_defs.append(
                py_fn_node(
                    name=arity_name,
                    args=ast.arguments(
                        args=fn_args,
                        kwarg=None,
                        vararg=varg,
                        kwonlyargs=[],
                        defaults=[],
                        kw_defaults=[],
                    ),
                    body=fn_body_ast,
                    decorator_list=[_TRAMPOLINE_FN_NAME]
                    if ctx.recur_point.has_recur
                    else [],
                    returns=None,
                )
            )

    dispatch_fn_ast = __multi_arity_dispatch_fn(
        ctx,
        py_fn_name,
        arity_to_name,
        default_name=rest_arity_name,
        max_fixed_arity=node.max_fixed_arity,
        meta_node=meta_node,
        is_async=node.is_async,
    )

    return GeneratedPyAST(
        node=dispatch_fn_ast.node,
        dependencies=list(chain(fn_defs, dispatch_fn_ast.dependencies)),
    )
def __multi_arity_fn_to_py_ast ( # pylint: disable=too-many-locals ctx : GeneratorContext , node : Fn , methods : Collection [ FnMethod ] , def_name : Optional [ str ] = None , meta_node : Optional [ MetaNode ] = None , ) -> GeneratedPyAST : assert node . op == NodeOp . FN assert all ( [ method . op == NodeOp . FN_METHOD for method in methods ] ) lisp_fn_name = node . local . name if node . local is not None else None py_fn_name = __fn_name ( lisp_fn_name ) if def_name is None else munge ( def_name ) py_fn_node = ast . AsyncFunctionDef if node . is_async else ast . FunctionDef arity_to_name = { } rest_arity_name : Optional [ str ] = None fn_defs = [ ] for method in methods : arity_name = f"{py_fn_name}__arity{'_rest' if method.is_variadic else method.fixed_arity}" if method . is_variadic : rest_arity_name = arity_name else : arity_to_name [ method . fixed_arity ] = arity_name with ctx . new_symbol_table ( arity_name ) , ctx . new_recur_point ( method . loop_id , RecurType . FN , is_variadic = node . is_variadic ) : # Allow named anonymous functions to recursively call themselves if lisp_fn_name is not None : ctx . symbol_table . new_symbol ( sym . symbol ( lisp_fn_name ) , py_fn_name , LocalType . FN ) fn_args , varg , fn_body_ast = __fn_args_to_py_ast ( ctx , method . params , method . body ) fn_defs . append ( py_fn_node ( name = arity_name , args = ast . arguments ( args = fn_args , kwarg = None , vararg = varg , kwonlyargs = [ ] , defaults = [ ] , kw_defaults = [ ] , ) , body = fn_body_ast , decorator_list = [ _TRAMPOLINE_FN_NAME ] if ctx . recur_point . has_recur else [ ] , returns = None , ) ) dispatch_fn_ast = __multi_arity_dispatch_fn ( ctx , py_fn_name , arity_to_name , default_name = rest_arity_name , max_fixed_arity = node . max_fixed_arity , meta_node = meta_node , is_async = node . is_async , ) return GeneratedPyAST ( node = dispatch_fn_ast . node , dependencies = list ( chain ( fn_defs , dispatch_fn_ast . dependencies ) ) , )
Return a Python AST node for a function with multiple arities.
Return a Python AST node for a function with multiple arities .