From 8fb2f65d4b1b391cfbed86cabd34837d8e8acc7a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?= =?UTF-8?q?=E0=A4=A4=29?= Date: Mon, 9 May 2022 20:01:14 +0545 Subject: [PATCH] TqdmCallback: make it lazy --- dvc/fs/_callback.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/dvc/fs/_callback.py b/dvc/fs/_callback.py index cbdb2adcb3..87e9101b40 100644 --- a/dvc/fs/_callback.py +++ b/dvc/fs/_callback.py @@ -1,7 +1,9 @@ +from contextlib import ExitStack from functools import wraps from typing import IO, TYPE_CHECKING, Any, Dict, Optional, TypeVar, cast import fsspec +from funcy import cached_property if TYPE_CHECKING: from typing import Callable @@ -96,17 +98,24 @@ class NoOpCallback(FsspecCallback, fsspec.callbacks.NoOpCallback): class TqdmCallback(FsspecCallback): def __init__(self, progress_bar: "Tqdm" = None, **tqdm_kwargs): + self._tqdm_kwargs = tqdm_kwargs + self._progress_bar = progress_bar + self._stack = ExitStack() + super().__init__() + + @cached_property + def progress_bar(self): from dvc.progress import Tqdm - self.progress_bar = progress_bar or Tqdm(**tqdm_kwargs) - super().__init__() + return self._stack.enter_context( + self._progress_bar or Tqdm(**self._tqdm_kwargs) + ) def __enter__(self): - self.progress_bar.__enter__() return self def close(self): - self.progress_bar.close() + self._stack.close() def set_size(self, size): if size is not None: