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:
    """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)),
    )
Return a Python AST node for a function with multiple arities.