From a7cbc95cc24b03fdceb90956e487ea03ccae4933 Mon Sep 17 00:00:00 2001
From: James Graham <J.A.Graham@soton.ac.uk>
Date: Tue, 12 Jul 2016 15:20:55 +0100
Subject: [PATCH] Add tqdm progress bar if available Specify which trajectory
 reader is being used

---
 pycgtool/frame.py     |  1 -
 pycgtool/interface.py | 18 +++++++++++++++++-
 pycgtool/pycgtool.py  |  2 +-
 test/test_frame.py    |  8 +++++---
 4 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/pycgtool/frame.py b/pycgtool/frame.py
index a7f7c97..dc4907a 100644
--- a/pycgtool/frame.py
+++ b/pycgtool/frame.py
@@ -124,7 +124,6 @@ class Frame:
     """
     Hold Atom data separated into Residues
     """
-
     def __init__(self, gro=None, xtc=None, itp=None, frame_start=0, xtc_reader="simpletraj"):
         """
         Return Frame instance having read Residues and Atoms from GRO if provided
diff --git a/pycgtool/interface.py b/pycgtool/interface.py
index e878d02..210a486 100644
--- a/pycgtool/interface.py
+++ b/pycgtool/interface.py
@@ -6,6 +6,11 @@ import curses
 import curses.textpad
 import time
 
+try:
+    from tqdm import tqdm
+except ImportError:
+    pass
+
 
 class Options:
     """
@@ -267,8 +272,19 @@ class Progress:
     def run(self):
         """
         Iterate through self until stopped by maximum iterations or False condition.
+
+        Use the tqdm library if it is present.
         """
-        collections.deque(self, maxlen=0)
+        if self._quiet:
+            collections.deque(self, maxlen=0)
+        else:
+            try:
+                self._quiet = True
+                collections.deque(tqdm(self, total=len(self)-1), maxlen=0)
+            except NameError:
+                self._quiet = False
+                collections.deque(self, maxlen=0)
+
         return self._its
 
     @property
diff --git a/pycgtool/pycgtool.py b/pycgtool/pycgtool.py
index 4696f6b..2660c1d 100755
--- a/pycgtool/pycgtool.py
+++ b/pycgtool/pycgtool.py
@@ -18,7 +18,7 @@ def main(args, config):
     :param args: Arguments from argparse
     :param config: Configuration dictionary
     """
-    frame = Frame(gro=args.gro, xtc=args.xtc, itp=args.itp, frame_start=args.begin, xtc_reader="mdtraj")
+    frame = Frame(gro=args.gro, xtc=args.xtc, itp=args.itp, frame_start=args.begin)
 
     if args.bnd:
         bonds = BondSet(args.bnd, config)
diff --git a/test/test_frame.py b/test/test_frame.py
index 59942f1..848e299 100644
--- a/test/test_frame.py
+++ b/test/test_frame.py
@@ -72,11 +72,13 @@ class FrameTest(unittest.TestCase):
         os.remove("water-out.gro")
 
     def test_frame_read_xtc_simpletraj_numframes(self):
-        frame = Frame(gro="test/data/water.gro", xtc="test/data/water.xtc")
+        frame = Frame(gro="test/data/water.gro", xtc="test/data/water.xtc",
+                      xtc_reader="simpletraj")
         self.assertEqual(12, frame.numframes)
 
-    def test_frame_read_xtc(self):
-        frame = Frame(gro="test/data/water.gro", xtc="test/data/water.xtc")
+    def test_frame_read_xtc_simpletraj(self):
+        frame = Frame(gro="test/data/water.gro", xtc="test/data/water.xtc",
+                      xtc_reader="simpletraj")
         self.assertEqual(663, frame.natoms)
         # These are the coordinates from the gro file
         np.testing.assert_allclose(np.array([0.696, 1.33, 1.211]),
-- 
GitLab