Loading big tables now shows a progress indicator. #180

pokedex/db/load.py

@@ -45,15 +45,22 @@ def _wildcards_to_regex(strings):
 
 
 def _get_verbose_prints(verbose):
-    """If `verbose` is true, returns two functions: one for printing a starting
-    message, and the other for printing a success or failure message when
-    finished.
+    """If `verbose` is true, returns three functions: one for printing a
+    starting message, one for printing an interim status update, and one for
+    printing a success or failure message when finished.
 
-    If `verbose` is false, returns two no-op functions.
+    If `verbose` is false, returns no-op functions.
     """
 
-    if verbose:
-        import sys
+    if not verbose:
+        # Return dummies
+        def dummy(*args, **kwargs):
+            pass
+
+        return dummy, dummy, dummy
+
+    ### Okay, verbose == True; print stuff
+
     def print_start(thing):
         # Truncate to 66 characters, leaving 10 characters for a success
         # or failure message
@@ -65,17 +72,39 @@ def _get_verbose_prints(verbose):
         print "%s...%s" % (truncated_thing, ' ' * num_spaces),
         sys.stdout.flush()
 
-        def print_done(msg='ok'):
-            print msg
+    if sys.stdout.isatty():
+        # stdout is a terminal; stupid backspace tricks are OK.
+        # Don't use print, because it always adds magical spaces, which
+        # makes backspace accounting harder
+
+        backspaces = [0]
+        def print_status(msg):
+            # Overwrite any status text with spaces before printing
+            sys.stdout.write('\b' * backspaces[0])
+            sys.stdout.write(' ' * backspaces[0])
+            sys.stdout.write('\b' * backspaces[0])
+            sys.stdout.write(msg)
             sys.stdout.flush()
+            backspaces[0] = len(msg)
 
-        return print_start, print_done
+        def print_done(msg='ok'):
+            # Overwrite any status text with spaces before printing
+            sys.stdout.write('\b' * backspaces[0])
+            sys.stdout.write(' ' * backspaces[0])
+            sys.stdout.write('\b' * backspaces[0])
+            sys.stdout.write(msg + "\n")
+            sys.stdout.flush()
+            backspaces[0] = 0
 
-    # Not verbose; return dummies
-    def dummy(*args, **kwargs):
+    else:
+        # stdout is a file (or something); don't bother with status at all
+        def print_status(msg):
             pass
 
-    return dummy, dummy
+        def print_done(msg='ok'):
+            print msg
+
+    return print_start, print_status, print_done
 
 
 def load(session, tables=[], directory=None, drop_tables=False, verbose=False):
@@ -101,7 +130,7 @@ def load(session, tables=[], directory=None, drop_tables=False, verbose=False):
     """
 
     # First take care of verbosity
-    print_start, print_done = _get_verbose_prints(verbose)
+    print_start, print_status, print_done = _get_verbose_prints(verbose)
 
 
     if not directory:
@@ -136,12 +165,15 @@ def load(session, tables=[], directory=None, drop_tables=False, verbose=False):
         print_start(table_name)
 
         try:
-            csvfile = open("%s/%s.csv" % (directory, table_name), 'rb')
+            csvpath = "%s/%s.csv" % (directory, table_name)
+            csvfile = open(csvpath, 'rb')
         except IOError:
             # File doesn't exist; don't load anything!
             print_done('missing?')
             continue
 
+        csvsize = os.stat(csvpath).st_size
+
         reader = csv.reader(csvfile, lineterminator='\n')
         column_names = [unicode(column) for column in reader.next()]
 
@@ -164,6 +196,9 @@ def load(session, tables=[], directory=None, drop_tables=False, verbose=False):
             session.commit()
             new_rows[:] = []
 
+            progress = "{0}%".format(100 * csvfile.tell() // csvsize)
+            print_status(progress)
+
         for csvs in reader:
             row_data = {}
 
@@ -254,7 +289,7 @@ def dump(session, tables=[], directory=None, verbose=False):
     """
 
     # First take care of verbosity
-    print_start, print_done = _get_verbose_prints(verbose)
+    print_start, print_status, print_done = _get_verbose_prints(verbose)
 
 
     if not directory: