Function wrap_as_closure

python/paddle/jit/dy2static/utils.py:479–571

Wrap a function to a closure function. Before: >>> def fn(x): ... After: >>> def create_fn(): ... closure_var_1 = None ... ... def fn(x): ... ... ... return fn ... ... ... fn = create_fn()

(tree: gast.AST, closure_vars: list[str])

def wrap_as_closure(tree: gast.AST, closure_vars: list[str]) -> gast.AST:
    """
    Wrap a function to a closure function.

    Before:

        >>> def fn(x): ...

    After:

        >>> def create_fn():
        ...     closure_var_1 = None
        ...
        ...     def fn(x): ...
        ...
        ...     return fn
        ...
        ...
        ... fn = create_fn()
    """

    def create_assign_node(name, value) -> gast.Assign:
        return gast.Assign(
            targets=[
                gast.Name(
                    id=name,
                    ctx=gast.Store(),
                    annotation=[],
                    type_comment=[],
                )
            ],
            value=value,
            type_comment=None,
        )

    def create_wrppper_fn_def_node(name, body) -> gast.FunctionDef:
        return gast.FunctionDef(
            name=name,
            args=gast.arguments(
                args=[],
                posonlyargs=[],
                vararg=None,
                kwonlyargs=[],
                kw_defaults=[],
                kwarg=None,
                defaults=[],
            ),
            body=body,
            decorator_list=[],
            returns=None,
            type_comment=None,
            type_params=[],
        )

    if not isinstance(tree, gast.Module):
        return tree
    if len(tree.body) != 1:
        return tree