From e8db31d066adf84b3def5d0bbfa0c5f3720e0135 Mon Sep 17 00:00:00 2001
From: Ilya Kreymer <ikreymer@gmail.com>
Date: Sun, 22 Mar 2015 21:50:56 -0700
Subject: [PATCH] cli: improve wayback cli to take optional port, threads and
 working dir arguments switch to waitress as default WSGI server instead of
 wsgiref

---
 .gitattributes           |  3 +++
 pywb/apps/cli.py         | 32 ++++++++++++++++++++++++++++++++
 pywb/apps/wayback.py     | 10 ++--------
 pywb/manager/manager.py  | 10 +++++++---
 setup.py                 |  3 ++-
 tests/test_auto_colls.py | 11 +++++++++++
 6 files changed, 57 insertions(+), 12 deletions(-)
 create mode 100644 pywb/apps/cli.py

diff --git a/.gitattributes b/.gitattributes
index 8e74ef52..3f4f86e3 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -1,4 +1,7 @@
 *.arc  -text
 *.warc -text
+*.idx  -text
+*.idxj -text
 *.cdx  -text
+*.cdxj -text
 *.gz   -text
diff --git a/pywb/apps/cli.py b/pywb/apps/cli.py
new file mode 100644
index 00000000..2972f8af
--- /dev/null
+++ b/pywb/apps/cli.py
@@ -0,0 +1,32 @@
+#=================================================================
+def wayback(args=None):
+    from argparse import ArgumentParser, RawTextHelpFormatter
+
+    parser = ArgumentParser('pywb Wayback Web Archive Replay')
+    parser.add_argument('-p', '--port', type=int, default=8080)
+    parser.add_argument('-t', '--threads', type=int, default=4)
+
+    help_dir='Specify root archive dir (default is current working directory)'
+    parser.add_argument('-d', '--directory', help=help_dir)
+
+    r = parser.parse_args(args)
+    if r.directory:  #pragma: no cover
+        import os
+        os.chdir(r.directory)
+
+    # Load App
+    from pywb.apps.wayback import application
+
+    try:
+        from waitress import serve
+        serve(application, port=r.port, threads=r.threads)
+    except ImportError:  # pragma: no cover
+        # Shouldn't ever happen as installing waitress, but just in case..
+        from pywb.framework.wsgi_wrappers import start_wsgi_server
+        start_wsgi_server(application, 'Wayback', default_port=r.port)
+
+
+#=================================================================
+if __name__ == "__main__":
+    wayback()
+
diff --git a/pywb/apps/wayback.py b/pywb/apps/wayback.py
index 17b7744f..4ba24d36 100644
--- a/pywb/apps/wayback.py
+++ b/pywb/apps/wayback.py
@@ -1,14 +1,8 @@
-from pywb.framework.wsgi_wrappers import init_app, start_wsgi_server
+from pywb.framework.wsgi_wrappers import init_app
 from pywb.webapp.pywb_init import create_wb_router
 
+
 #=================================================================
 # init pywb app
 #=================================================================
 application = init_app(create_wb_router, load_yaml=True)
-
-
-def main():  # pragma: no cover
-    start_wsgi_server(application, 'Wayback')
-
-if __name__ == "__main__":
-    main()
diff --git a/pywb/manager/manager.py b/pywb/manager/manager.py
index b40bb6e1..4a360c4d 100644
--- a/pywb/manager/manager.py
+++ b/pywb/manager/manager.py
@@ -55,6 +55,10 @@ directory structure expected by pywb
 
     def list_colls(self):
         print('Collections:')
+        if not os.path.isdir(self.colls_dir):
+            msg = ('"Collections" directory not found. ' +
+                   'To create a new collection, run:\n\n{0} init <name>')
+            raise IOError(msg.format(sys.argv[0]))
         for d in os.listdir(self.colls_dir):
             if os.path.isdir(os.path.join(self.colls_dir, d)):
                 print('- ' + d)
@@ -87,8 +91,9 @@ directory structure expected by pywb
 
     def _assert_coll_exists(self):
         if not os.path.isdir(self.curr_coll_dir):
-            raise IOError('Collection {0} does not exist'.
-                          format(self.coll_name))
+            msg = ('Collection {0} does not exist. ' +
+                   'To create a new collection, run\n\n{1} init {0}')
+            raise IOError(msg.format(self.coll_name, sys.argv[0]))
 
     def add_warcs(self, warcs):
         if not os.path.isdir(self.archive_dir):
@@ -466,7 +471,6 @@ def main_wrap_exc():  #pragma: no cover
     try:
         main()
     except Exception as e:
-        raise
         print('Error: ' + str(e))
         sys.exit(2)
 
diff --git a/setup.py b/setup.py
index f0a7c7f4..44c12365 100755
--- a/setup.py
+++ b/setup.py
@@ -75,6 +75,7 @@ setup(
         'surt',
         'pyyaml',
         'youtube_dl',
+        'waitress',
         'watchdog'
        ],
     tests_require=[
@@ -88,7 +89,7 @@ setup(
     test_suite='',
     entry_points="""
         [console_scripts]
-        wayback = pywb.apps.wayback:main
+        wayback = pywb.apps.cli:wayback
         cdx-server = pywb.apps.cdx_server:main
         cdx-indexer = pywb.warc.cdxindexer:main
         live-rewrite-server = pywb.apps.live_rewrite_server:main
diff --git a/tests/test_auto_colls.py b/tests/test_auto_colls.py
index 31e8fa89..029fe12d 100644
--- a/tests/test_auto_colls.py
+++ b/tests/test_auto_colls.py
@@ -70,6 +70,13 @@ class TestManagedColls(object):
     def _get_sample_warc(self, name):
         return os.path.join(get_test_dir(), 'warcs', name)
 
+    @patch('waitress.serve', lambda *args, **kwargs: None)
+    def test_run_cli(self):
+        """ test new wayback cli interface
+        """
+        from pywb.apps.cli import wayback
+        wayback([])
+
     def test_create_first_coll(self):
         """ Test first collection creation, with all required dirs
         """
@@ -561,6 +568,10 @@ class TestManagedColls(object):
 
         shutil.rmtree(colls)
 
+        # No Collections to list
+        with raises(IOError):
+            main(['list'])
+
         # No Collections
         self._create_app()
         resp = self.testapp.get('/test/', status=404)