def wrap_code(code: str, args: str = '') -> ast.Module:
    """
    Compiles Python code into an async function or generator,
    and automatically adds return if the function body is a single evaluation.
    Also adds inline import expression support.
    """

    if sys.version_info >= (3, 7):
        user_code = import_expression.parse(code, mode='exec')
        injected = ''
    else:
        injected = code

    mod = import_expression.parse(CORO_CODE.format(args, textwrap.indent(injected, ' ' * 8)), mode='exec')

    definition = mod.body[-1]  # async def ...:
    assert isinstance(definition, ast.AsyncFunctionDef)

    try_block = definition.body[-1]  # try:
    assert isinstance(try_block, ast.Try)

    if sys.version_info >= (3, 7):
        try_block.body.extend(user_code.body)
    else:
        ast.increment_lineno(mod, -16)  # bring line numbers back in sync with repl

    ast.fix_missing_locations(mod)

    is_asyncgen = any(isinstance(node, ast.Yield) for node in ast.walk(try_block))

    last_expr = try_block.body[-1]

    # if the last part isn't an expression, ignore it
    if not isinstance(last_expr, ast.Expr):
        return mod

    # if the last expression is not a yield
    if not isinstance(last_expr.value, ast.Yield):
        # copy the expression into a return/yield
        if is_asyncgen:
            # copy the value of the expression into a yield
            yield_stmt = ast.Yield(last_expr.value)
            ast.copy_location(yield_stmt, last_expr)
            # place the yield into its own expression
            yield_expr = ast.Expr(yield_stmt)
            ast.copy_location(yield_expr, last_expr)

            # place the yield where the original expression was
            try_block.body[-1] = yield_expr
        else:
            # copy the expression into a return
            return_stmt = ast.Return(last_expr.value)
            ast.copy_location(return_stmt, last_expr)

            # place the return where the original expression was
            try_block.body[-1] = return_stmt

    return mod
Compiles Python code into an async function or generator,
    and automatically adds return if the function body is a single evaluation.
    Also adds inline import expression support.