Capturing tqdm instances

tqdm_capture.py

import pytest
from funcy import wrap_with
import threading


class TdqmCapture:
    def __init__(self):
        self._instances = []

    @property
    def instances(self):
        return self._instances

    def __enter__(self):
        self.clear()
        yield self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def clear(self):
        self._instances.clear()

    @wrap_with(threading.Lock())
    def add(self, obj):
        self._instances.append(obj)
        return obj


@pytest.fixture
def captqdm(mocker):
    from dvc.progress import Tqdm

    tqdm_og_init = Tqdm.__init__
    result = TdqmCapture()

    def spy(obj, *args, **kwargs):
        return result.add(tqdm_og_init(obj, *args, **kwargs))

    mocker.patch('dvc.progress.Tqdm.__init__', spy)
    yield result