有没有办法给 PyYAML yaml.load() 函数添加一个进度条（比如 tqdm）？

Question

使用 PyYAML，CLoader 作为 YAML 解析器，我试图加载 YAML 文件，解析它，然后将它写入一个单独的文件.

出于测试目的，我使用了一个非常大的 YAML 文件，大于 1GB。

我正在尝试在命令行中包含一个进度条，以显示我的 Python 脚本是运行并估计需要多长时间。

这是我当前的代码：

import yaml
import argparse

from tqdm import tqdm
from yaml import CLoader as Loader

def main():

parser = argparse.ArgumentParser(description='Takes in YAML files and uploads straight to Neo4J database')
parser.add_argument('-f', '--files', nargs='+', metavar='', required=True,
                    help='<Required> One or more YAML files to upload')

args = parser.parse_args()

for file_name in args.files:

    with open(file_name, 'r') as stream:
        print("Reading input file...")
        with open('test2.txt', 'w') as wf:
            print("Writing to output file...")

            try:
                for data in tqdm(yaml.load(stream, Loader=Loader)):
                    wf.write(data.get('primaryName') + '\n')
                    wf.write('++++++++++\n')
            except yaml.YAMLError as exc:
                print(exc)

if __name__ == "__main__":
    main()

现在发生的情况是，数据写入循环显示 tqdm 进度条，但 yaml.load() 进程没有显示，该进程占用的时间最多。

即长时间不显示进度条，直到YAML文件加载完毕

我希望找到一个解决方案，使我能够在我无法访问的函数周围环绕一个进度条，在这种情况下，yaml.load()。

我做错了什么吗？任何建议都将不胜感激。

Answer 1

不，无法在您无权访问的代码周围设置进度条。

此外，当您循环遍历一个可迭代对象时，您只能使用 iterable-based 接口到 tqdm，而您不在此处。所以你必须使用基于 update 的界面：

with tqdm(total=100) as pbar:
    for i in range(10):
        pbar.update(10)

问题是，如何让 PyYAML 调用它 pbar.update？

理想情况下，您想找到一个地方来挂钩加载过程，您可以调用 pbar.update。如果那不可能，你将不得不做一些丑陋的事情（比如 fork PyYAML 并添加到它的 API，或者在运行时通过 monkeypatching 做同样的事情），或者切换到不同的图书馆。不过应该可以的。

显而易见的选择是创建您自己的 PyYAML.Loader 的子class。 PyYAML 的文档解释了此 class 的 API，因此您可以覆盖那里的任何方法以发出一些进度，然后 super 到基础 class.

但不幸的是 none 其中看起来很有前途。当然，您可以针对每个令牌、每个事件或每个节点调用一次，但是在不知道有多少令牌、事件或节点的情况下，这不会让您显示文件的深度。如果你想要一个不确定的进度微调器，那很好，但如果你能获得实际进度，并估计还需要多长时间等等，那就更好了。

你可以做的一件事是让你的 Loader subclass 调用 tell 上的 stream 来弄清楚到目前为止你读了多少字节。

我这台电脑上没有 PyYAML，文档也很混乱，所以你可能需要做一些实验，但它应该是这样的：

class ProgressLoader(yaml.CLoader):
    def __init__(self, stream, callback):
        super().__init__(stream)
        # __ because who knows what names the base class is using?
        self.__stream = stream
        self.__pos = 0
        self.__callback = callback
    def get_token(self):
        result = super().get_token()
        pos = self.__stream.tell()
        self.__callback(pos - self.__pos)
        self.__pos = pos
        return result

但是我不确定如何让 PyYAML 将你的回调传递给 ProgressLoader 构造函数，所以你必须做这样的事情：

with open(file_name, 'r') as stream:
    size = os.stat(stream.fileno()).st_size
    with tqdm(total=size) as progress:
        factory = lambda stream: ProgressLoader(stream, progress.update)
        data = yaml.load(stream, Loader=factory)

但是一旦我们无论如何都要访问文件，可能更容易不去纠结那些令人困惑的记录加载器类型，而只是编写一个文件包装器。

The docs for file objects 相当密集，但至少它们很清晰——实际工作非常简单：

class ProgressFileWrapper(io.TextIOBase):
    def __init__(self, file, callback):
        self.file = file
        self.callback = callback
    def read(self, size=-1):
        buf = self.file.read(size)
        if buf:
            self.callback(len(buf))
        return buf
    def readline(self, size=-1):
        buf = self.file.readline(size)
        if buf:
            self.callback(len(buf))
        return buf

现在：

with open(file_name, 'r') as stream:
    size = os.stat(stream.fileno()).st_size
    with tqdm(total=size) as progress:
        wrapper = ProgressFileWrapper(stream, progress.update)
        data = yaml.load(wrapper, Loader=Loader)

当然这并不完美。我们在这里假设所有的工作都是从磁盘读取文件，而不是解析它。这可能足够接近真实，我们可以侥幸逃脱，但如果不是，您将看到其中一个进度条快速移动到几乎 100%，然后毫无用处地停留在那里很长时间。¹

_{1.这不仅非常烦人，而且还与 Windows 和其他 Microsoft 产品紧密相关，以至于他们可能会起诉您盗用了它们的外观和风格。 :)}

有没有办法给 PyYAML yaml.load() 函数添加一个进度条（比如 tqdm）？

Is there a way to add a progress bar (such as tqdm) to PyYAML yaml.load() function?

python

pyyaml

progress-bar

tqdm