mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-02 07:18:25 +01:00
TqdmCallback: make it lazy
This commit is contained in:
committed by
skshetry
parent
477f266dc5
commit
8fb2f65d4b
+14
-5
@@ -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):
|
||||
from dvc.progress import Tqdm
|
||||
|
||||
self.progress_bar = progress_bar or Tqdm(**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
|
||||
|
||||
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:
|
||||
|
||||
Reference in New Issue
Block a user