python源码解读--flynt,将老式的字符串格式化代码转换为f-string

1. flynt

flynt是一个神奇的项目,它可以将你代码中的使用 % 和 fomat进行字符串格式化的代码转换为f-string形式的格式化代码,这个过程中,flynt会读取你的代码,解析你的语法,并修改你的代码,github地址是:https://github.com/ikamensh/flynt。

假设你有一个名为demo.py的文件

name_str = '小明'
age_int = 14
string ="我叫{name} 今年{age}岁了".format(name=name_str, age=age_int)

使用flynt,执行如下命令

flynt demo.py

再次打开demo.py文件,内容已经被修改

name_str = '小明'
age_int = 14
string =f"我叫{name_str} 今年{age_int}岁了"

很神奇吧,源码被改变了。它的实现原理是什么呢?显然不可能是通过字符串解析然后进行替换这么简单,原因在于,算法会非常复杂,难度太大。

你可能想到读取每一行代码,然后通过查找关键字format来判断这行代码是不是在做字符串格式化,然后提取关键字和format方法里的参数,之后进行替换。很不错,这是一个看起来可行的方法,但我随便写一行代码,就能让这个想法破产。

".format(name=name_str, age=age_int)"

在这行代码里,你可以找到关键字format,但它并不是在做字符串格式化,你可以增加新的判断条件,但不论怎样,我总能构建出一个满足你判断条件但又不是字符串格式化的代码,这条路是行不通的。

想要实现对字符串格式化代码的准确捕捉,必须借助抽象语法树。

2. ast,抽象语法树

ast是Abstract Syntax Trees的缩写,即抽象语法树。ast是源代码语法结构的抽象表示,它以树状结构来表现编程语言的语法结构,如果你对树这种数据结构比较了解,将有助于你理解抽象语法树。

python内置了ast模块,模块的parse方法接收一个字符串类型的参数,返回源码的抽象语法树,下面是一个两个int相加的代码

1 + 4

接下来,安装一个第三方模块,以便更好的显示语法树

pip install astpretty

让我们来看一下它的抽象语法树

import ast
import astpretty

code = '1 + 4'
tree = ast.parse(code)
astpretty.pprint(tree.body[0])

程序输出结果

Expr(           
    lineno=1,
    col_offset=0,
    value=BinOp(
        lineno=1,
        col_offset=0,
        left=Num(lineno=1, col_offset=0, n=1),
        op=Add(),
        right=Num(lineno=1, col_offset=4, n=4),
    ),
)

尽管输出的抽象语法树不是我们所熟悉的代码,但由于其结构化的设计和良好的命名习惯,我们可以很容易的猜测出它的含义:

  1. Expr 是表达式expression 的缩写
  2. lineno 表示行数
  3. col_offset 表示列的偏移数
  4. left 表示操作符的左侧
  5. Num 表示数字类型
  6. op=Add() 表示相加操作

python的解释器在执行代码时,需要将源码编译成字节码,其过程如下:

源代码解析 --> 语法树 --> 抽象语法树(AST) --> 控制流程图 --> 字节码

抽象语法树是其中的必要环节,我们平时写代码,不会去考虑抽象语法树,但对于一些特殊操作,只能借助抽象语法树来完成,比如检查源码里变量的命名是否符合规范。

2.1 检查变量名称是否符合规范

尽管有相应的编码规范,但不是所有人都愿意遵守,现在,请你写一个python脚本,该脚本可以检查指定源码里变量的命名是否符合规范,假设规范只要求变量名称小写且长度大于2。

你不可能通过字符串解析找出所有的变量,但是抽象语法树可以帮你做到这一点。源码被转换为语法树以后,你源码里的每一行代码,每一个操作,都会在语法树里表示出来,你只需要遍历这颗语法树并找到它。

假设下面的字符串code是一段需要检查的代码

code = """
a = 9
count = 10
Name = 'python'
c, d = 9, 10
"""

我们需要实现一个类,这个类继承ast.NodeTransformer

class VariableCheck(ast.NodeTransformer):
    pass

遍历语法树,需要调用visit方法,可语法树上有很多节点,我们只关心赋值语句的节点,在语法树上,是一个类型为Assign的节点,因此,我们需要实现visit_Assign

import ast
import _ast

code = """
a = 9
count = 10
Name = 'python'
c, d = 9, 10
"""

class VariableCheck(ast.NodeTransformer):
    def visit_Assign(self, node):
        target = node.targets[0]
        if isinstance(target, _ast.Name):
            varibale = target.id  # 变量名
            print(varibale)
        elif isinstance(target, _ast.Tuple):
            varibales = [item.id for item in target.elts]
            print(varibales)

tree = ast.parse(code)
variable_check = VariableCheck()
variable_check.visit(tree)

遍历语法树,VariableCheck实现了visit_Assign方法,所有的赋值语句的语法树节点都会被处理,变量名称就存储在node.targets[0]的id属性中,你问我是怎么知道的,我是通过pycharm 对代码进行debug,通过查看variables区域查出来的。

2.2 解析format方法里的参数

name_str = '小明'
age_int = 14
string ="我叫{name} 今年{age}岁了".format(name=name_str, age=age_int)

flynt可以将上面代码里的字符串格式化方式修改为f-string模式,这就需要识别出这一行代码在执行format方法,此外,还需要提炼出format的参数,遍历语法树时,可以实现visit_Call方法,所有的方法调用节点都会被处理

import ast

code ='"我叫{name} 今年{age}岁了".format(name=name_str, age=age_int)'
tree = ast.parse(code)

class FormatTransformer(ast.NodeTransformer):
    def visit_Call(self, node: ast.Call):
        if not self.is_format_call(node):
            return node

        var_map = {kw.arg: kw.value for kw in node.keywords}
        print(var_map)


    def is_format_call(self, node):
        call_from_string = isinstance(node.func.value, ast.Str) or (
                isinstance(node.func.value, ast.Constant)
                and isinstance(node.func.value.value, str)
        )
        return (
                isinstance(node, ast.Call)
                and hasattr(node.func, "value")
                and call_from_string
                and node.func.attr == "format"
        )

ft = FormatTransformer()
ft.visit(tree)

程序输出结果

{'name': <_ast.Name object at 0x0000013EA5894B70>, 
'age': <_ast.Name object at 0x0000013EA5894BE0>}

is_format_call 方法,是我直接复制粘贴的flynt源码,这个方法用来判断这一次的方法调用是不是字符串的format方法,仅仅判断node.func.attr的值是否为format是不够的,因为用户自定义类也可以定义实现format方法,因此还要考虑node.func.value的值是什么,这个value就是方法的调用者,如果是字符串,那就必然是在做字符串格式化了。

3. 提取字符串里关键字

前面的准备工作已经很充分了,通过解析语法树,得到了字符串格式化时format的两个参数,接下来,需要对源码里需要被格式化的字符串进行拆解,提取出字符串里的关键字,对于字符串

"我叫{name} 今年{age}岁了"

我们需要提取出name 和 age, 这部分工作,可以考虑用正则表达式来做

import re

pattern = re.compile(r'({[^}]*})')
text = "我叫{name} 今年{age}岁了"
res = pattern.findall(text)
print(res)

程序输出结果

['{name}', '{age}']

虽然也可以提取关键字,但并不满足最终的效果,在flynt源码里,作者的思路是提取出关键字,同时对字符串进行分解,关键字的部分要被2.2小节中的format参数替换掉,构建出一个新的语法树。

flynt的实现思路是使用string模块的Formatter类

import string

text = "我叫{name} 今年{age}岁了"
result = string.Formatter().parse(text)
for item in result:
    print(item)

程序输出结果

('我叫', 'name', '', None)
(' 今年', 'age', '', None)
('岁了', None, None, None)

Formatter类的parse方法专门用于解析待格式化的字符串,并将其拆解为多个部分,元组中的第2个元素就是需要填充的关键字,按照作者的思路,继续完善visit_Call方法

class FormatTransformer(ast.NodeTransformer):
    def visit_Call(self, node: ast.Call):
        if not self.is_format_call(node):
            return node

        var_map = {kw.arg: kw.value for kw in node.keywords}
        text = node.func.value.s
        # 拆解待格式化字符串
        splits = string.Formatter().parse(text)
        new_segments = []
        
        # 普通字符串转化为ast.Str, 关键字转换为ast.FormattedValue
        for raw, var_name, fmt_str, conversion in splits:
            if raw:
                new_segments.append(ast.Str(s=raw))

            if var_name is None:
                continue

            ast_name = var_map.pop(var_name)
            format_value = ast.FormattedValue(value=ast_name, conversion=-1, format_spec=None)
            new_segments.append(format_value)
        
        # 返回ast.JoinedStr,将各个部分连接起来
        return ast.JoinedStr(new_segments)

源码字符串被string.Formatter().parse处理后,普通的文字部分,别转换为_ast.Str类型数据,关键字部分被转换为_ast.FormattedValue类型数据。visit_Call方法最后的返回值类型是ast.JoinedStr。

4. 从语法树到代码

import astor

ft = FormatTransformer()
join_str_node = ft.visit(tree)
new_code = astor.to_source(join_str_node)
print(new_code)

使用ast可以将源码转换为抽象语法树,使用astor可以将抽象语法树转换为源码,这个库需要额外安装,程序输出结果

f"""我叫{name_str} 今年{age_int}岁了"""

将多出来的双引号替换成空字符串,就得到最终想要的结果了。

5. 总结

我并没有全面解读flynt库,而是重点讲解它是如何对抽象语法树进行遍历解析的,又是如何拼凑成一个新的语法树并使用第三方模块astor将其语法树解析成python代码。

源码里的实现,比我的实现要复杂的多,所考虑的情况更加全面,,我的代码不能用于生产,许多实现细节为了便于理解都忽略了,它只是比源码更容易理解而已。

扫描关注, 与我技术互动

QQ交流群: 211426309

加入知识星球, 每天收获更多精彩内容

分享日常研究的python技术和遇到的问题及解决方案