自动将python lambda重构为命名函数

Question

我正在从事purescript-python项目，有几个核心库广泛使用lambdas。由于代码的编译方式，lambda的位置最终变得模糊，每当出现错误时，就会产生奇怪的控制台消息

我希望重构这些库以尽可能少地使用lambdas。例如，如果有如下情况：

def foo(a):
  return lambda b: lambda c: lambda d: lambda e: a + b + c + d + e

这将是很好的生成

def foo(a):
  def _foo_internal_anon_1(b):
    def _foo_internal_anon_2(c):
        def _foo_internal_anon_3(d):
          def _foo_internal_anon_4(e):
            return a + b + c + d + e
          return _foo_internal_anon_4
        return _foo_internal_anon_3
    return _foo_internal_anon_2
  return _foo_internal_anon_1

有没有办法做到这一点，比如使用pylint或vscode或pycharm插件，还是必须手工完成

Answer 1

可以使用自定义^{}将Lambda内的Return节点转换为完整的函数定义。然后，可以借助CPython repo中的^{} tool来解析转换后的AST（从python3.9开始，您还可以使用^{}）。这允许转换整个脚本，而不仅仅是单个函数

这是节点转换器：

import ast
from contextlib import contextmanager


@contextmanager
def resetattr(obj, name, value):
    old_value = getattr(obj, name)
    setattr(obj, name, value)
    yield
    setattr(obj, name, old_value)


class ConvertLambda(ast.NodeTransformer):
    def __init__(self):
        super().__init__()
        self.base_name = None
        self.n = 0

    def visit_FunctionDef(self, node):
        if isinstance(node.body[-1], ast.Return) and isinstance(node.body[-1].value, ast.Lambda):
            lambda_node = node.body[-1].value
            with resetattr(self, 'base_name', self.base_name or node.name):
                with resetattr(self, 'n', self.n+1):
                    func_name = f'_{self.base_name}_internal_anon_{self.n}'
                    func_def = ast.FunctionDef(
                        name=func_name,
                        args=lambda_node.args,
                        body=[ast.Return(value=lambda_node.body)],
                        decorator_list=[],
                        returns=None,
                    )
                    self.visit(func_def)
            node.body.insert(-1, func_def)
            node.body[-1].value = ast.Name(id=func_name)
        return node

它可以与^{} class一起使用，如下所示（或者对于Python 3.9+，也可以与^{}一起使用）：

from unparse import Unparser

def convert_func_def(text):
    tree = ast.parse(text)
    tree = ast.fix_missing_locations(ConvertLambda().visit(tree))
    Unparser(tree)

默认情况下，这会将结果打印到sys.stdout，但是Unparser可以配置为使用任何文件，如对象：Unparser(tree, file=...)

这是针对示例函数获得的结果：

def foo(a):

    def _foo_internal_anon_1(b):

        def _foo_internal_anon_2(c):

            def _foo_internal_anon_3(d):

                def _foo_internal_anon_4(e):
                    return ((((a + b) + c) + d) + e)
                return _foo_internal_anon_4
            return _foo_internal_anon_3
        return _foo_internal_anon_2
    return _foo_internal_anon_1

它在添加的内容周围添加了一些额外的空行和括号，但也可以通过修改Unparser类来定制

Answer 2

不是完全自动的，但是可以使用^{}替换函数中的各种lambda。也许您可以将其转换为您喜爱的IDE的宏，该宏允许您突出显示一些文本，然后在其上运行转换

import functools
import re
import textwrap

def convert_func_def(text):
    """Assumes a multiple of 4 spaces as indentation."""
    func_name = re.search('def (.+?)(?=[(])', text).group(1)
    n_lambda = text.count('lambda')
    for i in range(n_lambda):
        text = re.sub(
            '^( +)return lambda (.+?): (.+?)$',
            functools.partial(replace, func_name=f'_{func_name}_internal_anon_{i+1}'),
            text,
            count=1,
            flags=re.MULTILINE,
        )
    return text

def replace(match, *, func_name):
    indent, args, body = match.groups()
    template = textwrap.dedent(f'''
        def {func_name}({args}):
            return {body}
        return {func_name}
    ''').strip()
    return textwrap.indent(template, indent)

应用于示例函数：

print(convert_func_def('''def foo(a):
    return lambda b: lambda c: lambda d: lambda e: a + b + c + d + e'''))

这是输出：

def foo(a):
    def _foo_internal_anon_1(b):
        def _foo_internal_anon_2(c):
            def _foo_internal_anon_3(d):
                def _foo_internal_anon_4(e):
                    return a + b + c + d + e
                return _foo_internal_anon_4
            return _foo_internal_anon_3
        return _foo_internal_anon_2
    return _foo_internal_anon_1