manager: templates: add collections manager (#74) commands for adding, removing and listing

available ui templates. Support for both collection and shared templates.
confirmation for overwrite/remove
updated full template list in default_config and added tests

pywb/default_config.yaml

@@ -14,12 +14,19 @@ paths:
 
     template_files:
         banner_html: banner.html
-        head_insert: head_insert.html
+        head_insert_html: head_insert.html
-        frame_insert: frame_insert.html
+        frame_insert_html: frame_insert.html
 
         query_html: query.html
         search_html: search.html
+        not_found_html: not_found.html
 
+    shared_template_files:
+        home_html: index.html
+        error_html: error.html
+
+        proxy_cert_download_html: proxy_cert_download.html
+        proxy_select_html: proxy_select.html
 
 head_insert_html: ui/head_insert.html
 frame_insert_html: ui/frame_insert.html

pywb/manager/manager.py

@@ -6,12 +6,22 @@ import logging
 from pywb.utils.loaders import load_yaml_config
 from pywb.utils.timeutils import timestamp20_now
 from pywb.warc.cdxindexer import main as cdxindexer_main
+from pywb.webapp.pywb_init import DEFAULT_CONFIG
+
+from distutils.util import strtobool
+from pkg_resources import resource_string
 
 from argparse import ArgumentParser, RawTextHelpFormatter
 import heapq
 import yaml
 
 
+#=============================================================================
+# to allow testing by mocking get_input
+def get_input(msg):  #pragma: no cover
+    return raw_input(msg)
+
+
 #=============================================================================
 class CollectionsManager(object):
     """ This utility is designed to
@@ -147,11 +157,113 @@ directory structure expected by pywb
             if len(v) != 2:
                 raise ValueError(msg)
 
+            print('Set {0}={1}'.format(v[0], v[1]))
             metadata[v[0]] = v[1]
 
         with open(metadata_yaml, 'w+b') as fh:
             fh.write(yaml.dump(metadata, default_flow_style=False))
 
+    def _load_templates_map(self):
+        defaults = load_yaml_config(DEFAULT_CONFIG)
+
+        # Coll Templates
+        templates = defaults['paths']['template_files']
+
+        for name, _ in templates.iteritems():
+            templates[name] = defaults[name]
+
+
+        # Shared Templates
+        shared_templates = defaults['paths']['shared_template_files']
+
+        for name, _ in shared_templates.iteritems():
+            shared_templates[name] = defaults[name]
+
+        return templates, shared_templates
+
+    def list_templates(self):
+        templates, shared_templates = self._load_templates_map()
+
+        print('Shared Templates')
+        for n, v in shared_templates.iteritems():
+            print('- {0}: (pywb/{1})'.format(n, v))
+
+        print('')
+
+        print('Collection Templates')
+        for n, v in templates.iteritems():
+            print('- {0}: (pywb/{1})'.format(n, v))
+
+    def _confirm_overwrite(self, full_path, msg):
+        if not os.path.isfile(full_path):
+            return True
+
+        res = get_input(msg)
+        try:
+            res = strtobool(res)
+        except ValueError:
+            res = False
+
+        if not res:
+            raise IOError('Skipping, {0} already exists'.format(full_path))
+
+    def _get_template_path(self, template_name, verb):
+        templates, shared_templates = self._load_templates_map()
+
+        try:
+            filename = templates[template_name]
+            if not self.coll_name:
+                msg = ('To {1} a "{0}" template, you must specify ' +
+                       'a collection name: template <coll> --{1} {0}')
+                raise IOError(msg.format(template_name, verb))
+
+            full_path = os.path.join(self.templates_dir,
+                                     os.path.basename(filename))
+
+        except KeyError:
+            try:
+                filename = shared_templates[template_name]
+                full_path = os.path.join(os.getcwd(),
+                                         os.path.basename(filename))
+            except KeyError:
+                msg = 'template name must be one of {0} or {1}'
+                msg.format(templates.keys(), shared_templates.keys())
+                raise KeyError(msg)
+
+        return full_path, filename
+
+    def add_template(self, template_name, force=False):
+        full_path, filename = self._get_template_path(template_name, 'add')
+
+        msg = ('Template file "{0}" ({1}) already exists. ' +
+               'Overwrite with default template? (y/n) ')
+        msg = msg.format(full_path, template_name)
+
+        if not force:
+            self._confirm_overwrite(full_path, msg)
+
+        data = resource_string('pywb', filename)
+        with open(full_path, 'w+b') as fh:
+            fh.write(data)
+
+        full_path = os.path.abspath(full_path)
+        print('Copied default template "{0}" to "{1}"'.format(filename, full_path))
+
+    def remove_template(self, template_name, force=False):
+        full_path, filename = self._get_template_path(template_name, 'remove')
+
+        if not os.path.isfile(full_path):
+            msg = 'Template "{0}" does not exist.'
+            raise IOError(msg.format(full_path))
+
+        msg = 'Delete template file "{0}" ({1})? (y/n) '
+        msg = msg.format(full_path, template_name)
+
+        if not force:
+            self._confirm_overwrite(full_path, msg)
+
+        os.remove(full_path)
+        print('Removed template file "{0}"'.format(full_path))
 
 #=============================================================================
 def main(args=None):
@@ -232,9 +344,37 @@ Create manage file based web archive collections
     metadata.add_argument('--set', nargs='+')
     metadata.set_defaults(func=do_metadata)
 
+    # Add default template
+    def do_add_template(r):
+        m = CollectionsManager(r.coll_name)
+        if r.add:
+            m.add_template(r.add, r.force)
+        elif r.remove:
+            m.remove_template(r.remove, r.force)
+        elif r.list:
+            m.list_templates()
+
+    template_help = 'Add default html template for customization'
+    template = subparsers.add_parser('template')
+    template.add_argument('coll_name', nargs='?', default='')
+    template.add_argument('-f', '--force', action='store_true')
+    template.add_argument('--add')
+    template.add_argument('--remove')
+    template.add_argument('--list', action='store_true')
+    template.set_defaults(func=do_add_template)
+
     r = parser.parse_args(args=args)
     r.func(r)
 
 
+# special wrapper for cli to avoid printing stack trace
+def main_wrap_exc():  #pragma: no cover
+    try:
+        main()
+    except Exception as e:
+        print('Error: ' + str(e))
+        sys.exit(2)
+
+
 if __name__ == "__main__":
-    main()
+    main_wrap_exc()

setup.py

@@ -90,7 +90,7 @@ setup(
         cdx-indexer = pywb.warc.cdxindexer:main
         live-rewrite-server = pywb.apps.live_rewrite_server:main
         proxy-cert-auth = pywb.framework.certauth:main
-        wb-manager = pywb.manager.manager:main
+        wb-manager = pywb.manager.manager:main_wrap_exc
         """,
     classifiers=[
         'Development Status :: 4 - Beta',

tests/test_auto_colls.py

@@ -14,6 +14,7 @@ from pywb import get_test_dir
 from pywb.framework.wsgi_wrappers import init_app
 
 from pytest import raises
+from mock import patch
 
 
 #=============================================================================
@@ -39,6 +40,14 @@ def teardown_module():
     os.chdir(orig_cwd)
 
 
+#=============================================================================
+mock_input_value = ''
+
+def mock_raw_input(*args):
+    global mock_input_value
+    return mock_input_value
+
+
 #=============================================================================
 class TestManagedColls(object):
     def setup(self):
@@ -266,6 +275,59 @@ class TestManagedColls(object):
         resp = self.testapp.get('/test/20140103030321/http://example.com?example=1')
         assert resp.status_int == 200
 
+    def test_add_default_templates(self):
+        """ Test add default templates: shared, collection,
+        and overwrite collection template
+        """
+        # list
+        main(['template', 'foo', '--list'])
+
+        # Add shared template
+        main(['template', '--add', 'home_html'])
+        assert os.path.isfile(os.path.join(self.root_dir, 'index.html'))
+
+        # Add collection template
+        main(['template', 'foo', '--add', 'query_html'])
+        assert os.path.isfile(os.path.join(self.root_dir, 'collections', 'foo', 'templates', 'query.html'))
+
+        # overwrite -- force
+        main(['template', 'foo', '--add', 'query_html', '-f'])
+
+    @patch('pywb.manager.manager.get_input', lambda x: 'y')
+    def test_add_template_input_yes(self):
+        """ Test answer 'yes' to overwrite
+        """
+        mock_raw_input_value = 'y'
+        main(['template', 'foo', '--add', 'query_html'])
+
+
+    @patch('pywb.manager.manager.get_input', lambda x: 'n')
+    def test_add_template_input_no(self):
+        """ Test answer 'no' to overwrite
+        """
+        with raises(IOError):
+            main(['template', 'foo', '--add', 'query_html'])
+
+    @patch('pywb.manager.manager.get_input', lambda x: 'other')
+    def test_add_template_input_other(self):
+        """ Test answer 'other' to overwrite
+        """
+        with raises(IOError):
+            main(['template', 'foo', '--add', 'query_html'])
+
+    @patch('pywb.manager.manager.get_input', lambda x: 'no')
+    def test_remove_not_confirm(self):
+        """ Test answer 'no' to remove
+        """
+        # don't remove -- not confirmed
+        with raises(IOError):
+            main(['template', 'foo', '--remove', 'query_html'])
+
+    @patch('pywb.manager.manager.get_input', lambda x: 'yes')
+    def test_remove_confirm(self):
+        # remove -- confirm
+        main(['template', 'foo', '--remove', 'query_html'])
+
     def test_no_templates(self):
         """ Test removing templates dir, using default template again
         """
@@ -284,8 +346,11 @@ class TestManagedColls(object):
         orig_stdout = sys.stdout
         buff = BytesIO()
         sys.stdout = buff
-        main(['list'])
+
-        sys.stdout = orig_stdout
+        try:
+            main(['list'])
+        finally:
+            sys.stdout = orig_stdout
 
         output = sorted(buff.getvalue().splitlines())
         assert len(output) == 4
@@ -294,6 +359,23 @@ class TestManagedColls(object):
         assert '- nested' in output
         assert '- test' in output
 
+    def test_err_template_remove(self):
+        """ Test various error conditions for templates:
+        invalid template name, no collection for collection template
+        no template file found
+        """
+        # no such template
+        with raises(KeyError):
+            main(['template', 'foo', '--remove', 'blah_html'])
+
+        # collection needed
+        with raises(IOError):
+            main(['template', '--remove', 'query_html'])
+
+        # already removed
+        with raises(IOError):
+            main(['template', 'foo', '--remove', 'query_html'])
+
     def test_err_no_such_coll(self):
         """ Test error adding warc to non-existant collection
         """