Commit f3ade0c3 authored by Markus Scheidgen's avatar Markus Scheidgen
Browse files

Added a ETA utility for cli commands.

parent 6eff468e
......@@ -13,9 +13,8 @@
# limitations under the License.
import click
from sys import stdout
from nomad import processing as proc, search, datamodel, infrastructure
from nomad import processing as proc, search, datamodel, infrastructure, utils
from nomad.cli.cli import cli
......@@ -39,14 +38,10 @@ def index(dry):
print('indexing %d ...' % all_calcs)
def calc_generator():
count = 0
for calc in proc.Calc.objects():
count += 1
if count % 1000 == 0:
stdout.write('\r indexed %10d of %10d calcs' % (count, all_calcs))
stdout.flush()
yield datamodel.CalcWithMetadata(**calc.metadata)
with utils.ETA(all_calcs, ' index %10d or %10d calcs, ETA %s') as eta:
for calc in proc.Calc.objects():
eta.add()
yield datamodel.CalcWithMetadata(**calc.metadata)
if dry:
for _ in calc_generator():
......
......@@ -47,6 +47,8 @@ import time
import re
from werkzeug.exceptions import HTTPException
import hashlib
import sys
from datetime import timedelta
from nomad import config
......@@ -403,3 +405,31 @@ class SleepTimeBackoff:
time.sleep(self.current_time)
self.current_time *= 2
self.current_time = min(self.max_time, self.current_time)
class ETA:
def __init__(self, total: int, message: str, interval: int = 1000):
self.start = time.time()
self.total = total
self.count = 0
self.interval = interval
self.interval_count = 0
self.message = message
def add(self, amount: int = 1):
self.count += amount
interval_count = int(self.count / self.interval)
if interval_count > self.interval_count:
self.interval_count = interval_count
delta_t = time.time() - self.start
eta = delta_t * (self.total - self.count) / self.count
eta_str = str(timedelta(seconds=eta))
sys.stdout.write('\r' + (self.message % (self.count, self.total, eta_str)))
sys.stdout.flush()
def __enter__(self):
self.start = time.time()
return self
def __exit__(self, *args, **kwargs):
print('')
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment