create dvc.progress.Tqdm, migrate dvc.remote.base

This commit is contained in:
Casper da Costa-Luis
2019-08-10 14:15:16 +01:00
parent 1511fd6dcc
commit 0ebce9c1e1
2 changed files with 38 additions and 27 deletions
+33 -8
View File
@@ -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
View File
@@ -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()))