mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-04 23:09:44 +01:00
create dvc.progress.Tqdm, migrate dvc.remote.base
This commit is contained in:
+33
-8
@@ -1,25 +1,50 @@
|
||||
"""Manages progress bars for dvc repo."""
|
||||
|
||||
from __future__ import print_function
|
||||
from __future__ import unicode_literals
|
||||
import logging
|
||||
from tqdm import tqdm
|
||||
from threading import Lock
|
||||
|
||||
from dvc.utils.compat import str
|
||||
|
||||
import sys
|
||||
import threading
|
||||
class Tqdm(tqdm):
|
||||
"""
|
||||
maximum-compatibility tqdm-based progressbars
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
iterable=None,
|
||||
disable=logging.getLogger(
|
||||
__name__).getEffectiveLevel() >= logging.CRITICAL,
|
||||
ascii=None, # TODO: True?
|
||||
**kwargs):
|
||||
"""
|
||||
kwargs : anything accepted by `tqdm.tqdm()`
|
||||
"""
|
||||
super(Tqdm, self).__init__(
|
||||
iterable=iterable,
|
||||
disable=disable,
|
||||
ascii=ascii,
|
||||
**kwargs)
|
||||
# self.set_lock(Lock())
|
||||
|
||||
def update_desc(self, desc, n=1):
|
||||
"""
|
||||
Calls `set_description(desc)` and `update(n)`
|
||||
"""
|
||||
self.set_description(desc, refresh=False)
|
||||
self.update(n)
|
||||
|
||||
|
||||
class Progress(tqdm):
|
||||
"""
|
||||
Simple multi-target progress bar.
|
||||
TODO: remove this class.
|
||||
"""
|
||||
def __init__(self):
|
||||
super(Progress, self).__init__(
|
||||
total=0,
|
||||
disable=logging.getLogger(__name__).getEffectiveLevel() >= logging.CRITICAL)
|
||||
self.set_lock(threading.Lock())
|
||||
disable=logging.getLogger(
|
||||
__name__).getEffectiveLevel() >= logging.CRITICAL)
|
||||
self.set_lock(Lock())
|
||||
self._targets = {}
|
||||
self.clearln()
|
||||
|
||||
@@ -77,7 +102,7 @@ class Progress(tqdm):
|
||||
|
||||
|
||||
class ProgressCallback(object):
|
||||
# TODO: is this meant to be a thin wrapper for multi-progress, or just one bar?
|
||||
# TODO: remove this thin wrapper
|
||||
def __init__(self, total):
|
||||
self.total = total
|
||||
self.current = 0
|
||||
|
||||
+5
-19
@@ -19,7 +19,7 @@ from dvc.exceptions import (
|
||||
ConfirmRemoveError,
|
||||
DvcIgnoreInCollectedDirError,
|
||||
)
|
||||
from dvc.progress import progress, ProgressCallback
|
||||
from dvc.progress import Tqdm
|
||||
from dvc.utils import (
|
||||
LARGE_DIR_SIZE,
|
||||
tmp_fname,
|
||||
@@ -159,7 +159,7 @@ class RemoteBASE(object):
|
||||
"This is only done once."
|
||||
)
|
||||
logger.info(msg)
|
||||
tasks = progress(tasks, total=len(file_infos))
|
||||
tasks = Tqdm(tasks, total=len(file_infos), unit="md5")
|
||||
|
||||
checksums = {
|
||||
file_infos[index]: task for index, task in enumerate(tasks)
|
||||
@@ -444,9 +444,6 @@ class RemoteBASE(object):
|
||||
|
||||
name = name or from_info.name
|
||||
|
||||
if not no_progress_bar:
|
||||
progress.update_target(name, 0, None)
|
||||
|
||||
try:
|
||||
self._upload(
|
||||
from_info.fspath,
|
||||
@@ -459,9 +456,6 @@ class RemoteBASE(object):
|
||||
logger.exception(msg.format(from_info, to_info))
|
||||
return 1 # 1 fail
|
||||
|
||||
if not no_progress_bar:
|
||||
progress.finish_target(name)
|
||||
|
||||
return 0
|
||||
|
||||
def download(
|
||||
@@ -490,11 +484,6 @@ class RemoteBASE(object):
|
||||
|
||||
name = name or to_info.name
|
||||
|
||||
if not no_progress_bar:
|
||||
# real progress is not always available,
|
||||
# lets at least show start and finish
|
||||
progress.update_target(name, 0, None)
|
||||
|
||||
makedirs(to_info.parent, exist_ok=True, mode=dir_mode)
|
||||
tmp_file = tmp_fname(to_info)
|
||||
|
||||
@@ -509,9 +498,6 @@ class RemoteBASE(object):
|
||||
|
||||
move(tmp_file, to_info, mode=file_mode)
|
||||
|
||||
if not no_progress_bar:
|
||||
progress.finish_target(name)
|
||||
|
||||
return 0
|
||||
|
||||
def open(self, path_info, mode="r", encoding=None):
|
||||
@@ -642,10 +628,10 @@ class RemoteBASE(object):
|
||||
Returns:
|
||||
A list with checksums that were found in the remote
|
||||
"""
|
||||
progress_callback = ProgressCallback(len(checksums))
|
||||
pbar = Tqdm(total=len(checksums))
|
||||
|
||||
def exists_with_progress(chunks):
|
||||
return self.batch_exists(chunks, callback=progress_callback)
|
||||
return self.batch_exists(chunks, callback=pbar.update_desc)
|
||||
|
||||
if self.no_traverse and hasattr(self, "batch_exists"):
|
||||
with ThreadPoolExecutor(max_workers=jobs or self.JOBS) as executor:
|
||||
@@ -654,7 +640,7 @@ class RemoteBASE(object):
|
||||
results = executor.map(exists_with_progress, chunks)
|
||||
in_remote = itertools.chain.from_iterable(results)
|
||||
ret = list(itertools.compress(checksums, in_remote))
|
||||
progress_callback.finish("")
|
||||
pbar.close()
|
||||
return ret
|
||||
|
||||
return list(set(checksums) & set(self.all()))
|
||||
|
||||
Reference in New Issue
Block a user