From 35190cea6dda6d65537b110f82a5d3f5277bd94c Mon Sep 17 00:00:00 2001
From: Oleg Broytman <phd@phdru.name>
Date: Sat, 24 Mar 2018 01:39:46 +0300
Subject: [PATCH] Fix(scripts/ml-import.py): Display progress bar on updating
 counters

---
 docs/news.rst        |  2 ++
 m_librarian/db.py    | 26 +++++++++++++++++++++++++-
 scripts/ml-import.py |  9 ++++++++-
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/docs/news.rst b/docs/news.rst
index 03347a1..618a282 100644
--- a/docs/news.rst
+++ b/docs/news.rst
@@ -4,6 +4,8 @@ News
 Version 0.0.17 (2018-03-??)
 ---------------------------
 
+* Display progress bar on updating counters.
+
 * Inhibit progress bar if stdout is not terminal.
 
 Version 0.0.16 (2018-03-22)
diff --git a/m_librarian/db.py b/m_librarian/db.py
index d81619d..9ae5ce0 100755
--- a/m_librarian/db.py
+++ b/m_librarian/db.py
@@ -206,18 +206,42 @@ def insert_author(surname, name, misc_name):
         return Author(surname=surname, name=name, misc_name=misc_name, count=0)
 
 
-def update_counters():
+def update_counters(pbar_cb=None):
+    if pbar_cb:
+        count = 0
+        for table in Author, Extension, Genre, Language:
+            count += table.select().count()
+        pbar_cb.set_max(count)
+
+    if pbar_cb:
+        count = 0
+
     for author in Author.select():
         author.count = AuthorBook.select(AuthorBook.q.author == author).count()
+        if pbar_cb:
+            count += 1
+            pbar_cb.display(count)
 
     for ext in Extension.select():
         ext.count = Book.select(Book.q.extension == ext.id).count()
+        if pbar_cb:
+            count += 1
+            pbar_cb.display(count)
 
     for genre in Genre.select():
         genre.count = BookGenre.select(BookGenre.q.genre == genre).count()
+        if pbar_cb:
+            count += 1
+            pbar_cb.display(count)
 
     for language in Language.select():
         language.count = Book.select(Book.q.language == language.id).count()
+        if pbar_cb:
+            count += 1
+            pbar_cb.display(count)
+
+    if pbar_cb:
+        pbar_cb.close()
 
 
 def test():
diff --git a/scripts/ml-import.py b/scripts/ml-import.py
index 169a715..07c7fbb 100755
--- a/scripts/ml-import.py
+++ b/scripts/ml-import.py
@@ -41,4 +41,11 @@ if __name__ == '__main__':
             print("Ok")
         else:
             import_inpx(inpx)
-    update_counters()
+    if use_pbar:
+        print("Updating counters", end=': ')
+        sys.stdout.flush()
+        pbar = ml_ttyProgressBar()
+        update_counters(pbar_cb=pbar)
+        print("Ok")
+    else:
+        update_counters()
-- 
2.39.2