cli refactor: use classes in cli to allow custom options

get rid of custom init for live_rewrite_handler, just use create_wb_router()
with custom config for consistent init

pywb/apps/cli.py

@@ -1,84 +1,126 @@
 import os
 import logging
 from argparse import ArgumentParser
+from pywb.framework.wsgi_wrappers import init_app
+from pywb.webapp.pywb_init import create_cdx_server_app, create_wb_router
 
 
 #=================================================================
 def cdx_server(args=None):  #pragma: no cover
-    def load():
+    CdxCli(args=args,
-        import pywb.apps.cdx_server as mod
+           default_port=8080,
-        return mod
+           desc='pywb CDX Index Server').run()
-
-    cli(appload=load,
-        args=args,
-        default_port=8090,
-        desc='pywb CDX Index Server')
 
 
 #=================================================================
 def live_rewrite_server(args=None):  #pragma: no cover
-    def load():
+    LiveCli(args=args,
-        import pywb.apps.live_rewrite_server as mod
+            default_port=8090,
-        return mod
+            desc='pywb Live Rewrite Proxy Server').run()
-
-    cli(appload=load,
-        args=args,
-        default_port=8090,
-        desc='pywb Live Rewrite Proxy Server')
 
 
 #=================================================================
 def wayback(args=None):
-    def load():
+    WaybackCli(args=args,
-        import pywb.apps.wayback as mod
+               default_port=8080,
-        return mod
+               desc='pywb Wayback Web Archive Replay').run()
-
-    cli(appload=load,
-        args=args,
-        default_port=8080,
-        desc='pywb Wayback Web Archive Replay')
 
 
-#=================================================================
+#=============================================================================
-def cli(appload, args=None, default_port=8080, desc=''):
+class BaseCli(object):
-    parser = ArgumentParser(desc)
+    def __init__(self, args=None, default_port=8080, desc=''):
-    parser.add_argument('-p', '--port', type=int, default=default_port)
+        parser = ArgumentParser(desc)
-    parser.add_argument('-t', '--threads', type=int, default=4)
+        parser.add_argument('-p', '--port', type=int, default=default_port)
-    parser.add_argument('-a', '--autoindex', action='store_true')
+        parser.add_argument('-t', '--threads', type=int, default=4)
 
-    help_dir='Specify root archive dir (default is current working directory)'
+        self.desc = desc
-    parser.add_argument('-d', '--directory', help=help_dir)
 
-    r = parser.parse_args(args)
+        self._extend_parser(parser)
-    if r.directory:  #pragma: no cover
-        os.chdir(r.directory)
 
-    # Load App
+        self.r = parser.parse_args(args)
-    #from pywb.apps.wayback import application
-    application = appload().application
 
-    if r.autoindex:
+        self.application = self.load()
-        from pywb.manager.manager import CollectionsManager
-        m = CollectionsManager('', must_exist=False)
-        if not os.path.isdir(m.colls_dir):
-            msg = 'No managed directory "{0}" for auto-indexing'
-            logging.error(msg.format(m.colls_dir))
-            import sys
-            sys.exit(2)
-        else:
-            msg = 'Auto-Indexing Enabled on "{0}"'
-            logging.info(msg.format(m.colls_dir))
-            m.autoindex(do_loop=False)
 
-    try:
+    def _extend_parser(self, parser):  #pragma: no cover
-        from waitress import serve
+        pass
-        serve(application, port=r.port, threads=r.threads)
+
-    except ImportError:  # pragma: no cover
+    def load(self):  #pragma: no cover
-        # Shouldn't ever happen as installing waitress, but just in case..
+        pass
-        from pywb.framework.wsgi_wrappers import start_wsgi_ref_server
+
-        start_wsgi_ref_server(application, desc, port=r.port)
+    def run(self):
+        try:
+            from waitress import serve
+            print(self.desc)
+            serve(self.application, port=self.r.port, threads=self.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_ref_server
+            start_wsgi_ref_server(self.application, self.desc, port=self.r.port)
 
 
-#=================================================================
+#=============================================================================
+class LiveCli(BaseCli):
+    def _extend_parser(self, parser):
+        parser.add_argument('-x', '--proxy',
+                            help='Specify host:port to use as HTTP/S proxy')
+
+        parser.add_argument('-f', '--framed', action='store_true',
+                            help='Replay using framed wrapping mode')
+
+    def load(self):
+        config = dict(proxyhostport=self.r.proxy,
+                      framed_replay=self.r.framed,
+                      enable_auto_colls=False,
+                      collections=dict(rewrite='$liveweb'))
+
+        return init_app(create_wb_router, load_yaml=False, config=config)
+
+
+#=============================================================================
+class ReplayCli(BaseCli):
+    def _extend_parser(self, parser):
+        parser.add_argument('-a', '--autoindex', action='store_true')
+
+        help_dir='Specify root archive dir (default is current working directory)'
+        parser.add_argument('-d', '--directory', help=help_dir)
+
+
+    def load(self):
+        if self.r.directory:  #pragma: no cover
+            os.chdir(self.r.directory)
+
+    def run(self):
+        if self.r.autoindex:
+            from pywb.manager.manager import CollectionsManager
+            m = CollectionsManager('', must_exist=False)
+            if not os.path.isdir(m.colls_dir):
+                msg = 'No managed directory "{0}" for auto-indexing'
+                logging.error(msg.format(m.colls_dir))
+                import sys
+                sys.exit(2)
+            else:
+                msg = 'Auto-Indexing Enabled on "{0}"'
+                logging.info(msg.format(m.colls_dir))
+                m.autoindex(do_loop=False)
+
+        super(ReplayCli, self).run()
+
+#=============================================================================
+class CdxCli(ReplayCli):  #pragma: no cover
+    def load(self):
+        super(CdxCli, self).load()
+        return init_app(create_cdx_server_app,
+                        load_yaml=True)
+
+
+#=============================================================================
+class WaybackCli(ReplayCli):
+    def load(self):
+        super(WaybackCli, self).load()
+        return init_app(create_wb_router,
+                        load_yaml=True)
+
+
+#=============================================================================
 if __name__ == "__main__":
     wayback()

pywb/apps/live_rewrite_server.py

@@ -1,34 +1,7 @@
-from pywb.framework.wsgi_wrappers import init_app
+from cli import LiveCli
-
-from pywb.webapp.live_rewrite_handler import create_live_rewriter_app
-
-from argparse import ArgumentParser
-
 
 #=================================================================
-# init rewrite server app
+# init default live rewrite server app
 #=================================================================
 
-def create_app():
+application = LiveCli([]).application
-    parser = ArgumentParser(description='Live Rewrite Server')
-
-    parser.add_argument('-x', '--proxy',
-                        action='store',
-                        help='Specify host:port to use as HTTP/S proxy')
-
-    parser.add_argument('-f', '--framed',
-                        action='store_true',
-                        help='Replay using framed wrapping mode')
-
-    result, unknown = parser.parse_known_args()
-
-    config = dict(proxyhostport=result.proxy,
-                  framed_replay=result.framed)
-
-    app = init_app(create_live_rewriter_app, load_yaml=False,
-                   config=config)
-
-    return app
-
-
-application = create_app()

pywb/webapp/live_rewrite_handler.py

@@ -1,6 +1,4 @@
-from pywb.framework.basehandlers import WbUrlHandler
 from pywb.framework.wbrequestresponse import WbResponse
-from pywb.framework.archivalrouter import ArchivalRouter, Route
 from pywb.framework.cache import create_cache
 
 from pywb.rewrite.rewrite_live import LiveRewriter
@@ -281,12 +279,3 @@ class YoutubeDLWrapper(object):
             self.htmlparser.locatestarttagend = self.orig_tagregex
 
         return info
-
-
-#=================================================================
-def create_live_rewriter_app(config={}):
-    routes = [Route('rewrite', RewriteHandler(config)),
-              Route('static/__pywb', StaticHandler('pywb/static/'))
-             ]
-
-    return ArchivalRouter(routes, hostpaths=['http://localhost:8080'])

tests/test_live_proxy.py

@@ -3,7 +3,8 @@ from BaseHTTPServer import HTTPServer, BaseHTTPRequestHandler
 
 from server_thread import ServerThreadRunner
 
-from pywb.webapp.live_rewrite_handler import create_live_rewriter_app, RewriteHandler
+from pywb.webapp.live_rewrite_handler import RewriteHandler
+from pywb.webapp.pywb_init import create_wb_router
 
 from pywb.framework.wsgi_wrappers import init_app
 import webtest
@@ -61,9 +62,13 @@ class TestProxyLiveRewriter:
 
         self.server = ServerThreadRunner(make_httpd)
 
-        self.app = init_app(create_live_rewriter_app, load_yaml=False,
+        config = dict(collections=dict(rewrite='$liveweb'),
-                           config=dict(framed_replay=True,
+                      framed_replay=True,
-                           proxyhostport=self.server.proxy_dict))
+                      proxyhostport=self.server.proxy_dict)
+
+        self.app = init_app(create_wb_router,
+                            load_yaml=False,
+                            config=config)
 
         def create_cache():
             return self.cache

tests/test_live_rewriter.py

@@ -1,11 +1,11 @@
-from pywb.webapp.live_rewrite_handler import create_live_rewriter_app, RewriteHandler
+from pywb.webapp.live_rewrite_handler import RewriteHandler
+from pywb.apps.cli import LiveCli
 from pywb.framework.wsgi_wrappers import init_app
 import webtest
 
 class TestLiveRewriter:
     def setup(self):
-        self.app = init_app(create_live_rewriter_app, load_yaml=False,
+        self.app = LiveCli(['-f']).application
-                            config=dict(framed_replay=True))
         self.testapp = webtest.TestApp(self.app)
 
     def test_live_rewrite_1(self):
@@ -28,7 +28,7 @@ class TestLiveRewriter:
         resp = self.testapp.get('/rewrite/tf_/http://example.com/')
         assert resp.status_int == 200
         assert '<iframe ' in resp.body
-        assert 'src="/rewrite/http://example.com/"' in resp.body
+        assert 'src="http://localhost:80/rewrite/http://example.com/"' in resp.body, resp.body
 
     def test_live_invalid(self):
         resp = self.testapp.get('/rewrite/http://abcdef', status=400)