python - 如何为已安装模块中发生的迭代创建进度条

Question

我的目标是为已安装模块内发生的迭代创建一个进度条。

为了在用户定义的函数中为迭代创建进度条，我传递了一个 tqdm.notebook.tqdm_notebook 对象作为可迭代对象:

import time
import numpy as np
from tqdm.notebook import tqdm

def iterate(over):
    for x in over: # creating progress bar for this
        print(x, end='')
        time.sleep(0.5)

xs = np.arange(5)
tqdm_xs = tqdm(xs) # creating tqdm.notebook.tqdm_notebook object

iterate(tqdm_xs) # progress bar, as expected
iterate(xs) # no progress bar

哪个有效:

但是，当我尝试对已安装模块内的 for 循环执行相同操作时，失败了。在 Astropy 的 Photutils 模块中，有一个 for label in labels 行 ( here )，我可以传递标签对象。

可重现的示例(主要基于 this - 安装 photutils 后工作:pip install photutils):

import photutils.datasets as phdat
import photutils.segmentation as phsegm
import astropy.convolution as conv
import astropy.stats as stats

data = phdat.make_100gaussians_image()
threshold = phsegm.detect_threshold(data, nsigma=2.)
sigma = 1.5
kernel = conv.Gaussian2DKernel(sigma, x_size=3, y_size=3)
kernel.normalize()
segm = phsegm.detect_sources(data, threshold, npixels=5, kernel=kernel)

这个有效:

segm_deblend = phsegm.deblend_sources(data, segm, npixels=5, kernel=kernel,
                                      nlevels=32, contrast=0.001, labels = segm.labels)

尝试传递 tqdm.notebook.tqdm_notebook 对象来创建进度条:

tqdm_segm_labels = tqdm(segm.labels)
segm_deblend = phsegm.deblend_sources(data, segm, npixels=5, kernel=kernel,
                                    nlevels=32, contrast=0.001, labels = tqdm_segm_labels)

我得到一个 AttributeError: 'int' object has no attribute '_comparable'。完整回溯:

0%
0/92 [00:00<?, ?it/s]
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-8-d101466650ae> in <module>()
      1 tqdm_segm_labels = tqdm(segm.labels)
      2 segm_deblend = phsegm.deblend_sources(data, segm, npixels=5, kernel=kernel,
----> 3                                     nlevels=32, contrast=0.001, labels = tqdm_segm_labels)

4 frames
/usr/local/lib/python3.7/dist-packages/astropy/utils/decorators.py in wrapper(*args, **kwargs)
    534                     warnings.warn(message, warning_type, stacklevel=2)
    535 
--> 536             return function(*args, **kwargs)
    537 
    538         return wrapper

/usr/local/lib/python3.7/dist-packages/photutils/segmentation/deblend.py in deblend_sources(data, segment_img, npixels, kernel, labels, nlevels, contrast, mode, connectivity, relabel)
    112         labels = segment_img.labels
    113     labels = np.atleast_1d(labels)
--> 114     segment_img.check_labels(labels)
    115 
    116     if kernel is not None:

/usr/local/lib/python3.7/dist-packages/photutils/segmentation/core.py in check_labels(self, labels)
    355 
    356         # check for positive label numbers
--> 357         idx = np.where(labels <= 0)[0]
    358         if idx.size > 0:
    359             bad_labels.update(labels[idx])

/usr/local/lib/python3.7/dist-packages/tqdm/utils.py in __le__(self, other)
     70 
     71     def __le__(self, other):
---> 72         return (self < other) or (self == other)
     73 
     74     def __eq__(self, other):

/usr/local/lib/python3.7/dist-packages/tqdm/utils.py in __lt__(self, other)
     67     """Assumes child has self._comparable attr/@property"""
     68     def __lt__(self, other):
---> 69         return self._comparable < other._comparable
     70 
     71     def __le__(self, other):

AttributeError: 'int' object has no attribute '_comparable'

解决方法是修改 Photutils 并在其中使用 tqdm(我在 this fork 上做了，它有效)，但这似乎有点矫枉过正，我希望有更简单的方法做这个。

Answer 1

当然一般没有办法直接修改一些不是你自己写的现有代码(是否“安装”不是问题)。

如果您认为它确实具有一般用途或兴趣，您可以提出一个补丁以允许此函数采用，例如，一个回调函数来调用每个循环。如果它通常是一个缓慢的函数，它可能会有用(我确实注意到实现中的一些事情可以改变以加快它的速度，但这是另一回事)。

您当然可以找到许多聪明的 hack 来使其在这个特定情况下工作，尽管考虑到它是专门针对此功能的实现细节设计的 hack，它会很脆弱。我发现了一些可能性。

最简单的似乎是这个愚蠢的把戏:

制作一个 ndarray 子类(我称之为 tqdm_array )，当在 Python 中迭代时，它会返回一个迭代器，该迭代器覆盖数组本身的 tqdm 进度条:

class tqdm_array(np.ndarray):
    def __iter__(self):
        return iter(tqdm.tqdm(np.asarray(self)))

然后准备打电话的时候deblend_sources把你的标签包在里面:

labels = np.array(segm_image.labels).view(tqdm_array)

并将其传递给 deblend_sources(..., labels=labels, ...) .

这会起作用，因为即使 labels由 NumPy 代码迭代，它将使用内部 C 代码直接迭代数组缓冲区(例如，对于像 labels <= 0 这样的操作。在大多数情况下，它不会调用 Python 级别的 __iter__ 方法，尽管可能有异常(exception)...

但是当遇到类似 for label in labels: 的 Python for 循环时(恰好在这个函数中只有一个)，你会得到你的进度条。