shot them all: use argparse to make it more generic

bin/shot_them_all.py

@@ -1,42 +1,85 @@
 # -*- coding: utf-8 -*-
 #
-# hobo-manage tenant_command runscript shot_them_all.py -d hobo.fred.local.0d.be
+# Take desktop/mobile shots of all themes from publik-base-theme.
+#
+# Usage:
+#  --directory DIRECTORY -- directory to store shots, defaults to shots/
+#  --no-headless         -- run in an actual browser window
+#  --no-memcached        -- do not restart memcached after theme changes
+#  --overlay OVERLAY     -- limit to given overlay ("none" to exclude overlays)
+#  --reshot              -- retake existing shots
+#  --timeout TIMEOUT     -- timeout between shots (time for hobo deploy)
+#  --theme THEME         -- theme to shot (substring match)
+#  URL                   -- specific URL to shot
+#
+# Example:
+#  hobo-manage tenant_command runscript shot_them_all.py -d hobo.fred.local.0d.be
 #
-# Take desktop/mobile shots of all themes from publik-base-theme, tailored for
-# my local installation.
 
+import argparse
 import os
 import json
 import time
 from selenium import webdriver
-from hobo.theme.utils import set_theme
+from hobo.theme.utils import get_selected_theme, get_themes, set_theme
 from hobo.deploy.signals import notify_agents
 import sys
 
-if not os.path.exists('shots'):
-    os.mkdir('shots')
+parser = argparse.ArgumentParser()
+parser.add_argument('--directory', dest='directory', type=str, default='shots')
+parser.add_argument('--no-headless', dest='no_headless', action='store_true')
+parser.add_argument('--no-memcached', dest='no_memcached', action='store_true')
+parser.add_argument('--overlay', dest='overlay', type=str, default='none')
+parser.add_argument('--reshot', dest='reshot', action='store_true')
+parser.add_argument('--timeout', dest='timeout', type=int, default=20)  # seconds
+parser.add_argument('--theme', dest='theme', type=str)
+parser.add_argument('url', metavar='URL', type=str, nargs='?',
+        default='https://auquo.fred.local.0d.be/contactez-nous/inscription-sur-les-listes/')
+parser.print_help()
+args = parser.parse_args()
+
+if not os.path.exists(args.directory):
+    os.mkdir(args.directory)
 
 options = webdriver.ChromeOptions()
-options.add_argument('headless')
+if not args.no_headless:
+    options.add_argument('headless')
 browser = webdriver.Chrome(chrome_options=options)
 
-themes_data = json.load(open('/home/fred/src/eo/publik-base-theme/themes.json'))
-for theme in themes_data['themes']:
+for theme in get_themes():
     theme_id = theme['id']
+    if args.overlay == 'all':
+        pass
+    elif args.overlay == 'none' and not theme.get('overlay'):
+        pass
+    elif args.overlay == theme.get('overlay'):
+        pass
+    else:
+        continue
+    if args.theme and not args.theme in theme['id']:
+        continue
     sys.stderr.write("%-25s" % theme_id)
-    set_theme(theme_id)
-    notify_agents(None)
-    for i in range(20):  # 20 seconds
-        sys.stderr.write('.')
-        time.sleep(1)
-    os.system('sudo service memcached restart')
-    browser.get('https://auquo.fred.local.0d.be/contactez-nous/inscription-sur-les-listes/')
-    browser.set_window_size(1400, 1000)
-    browser.save_screenshot('shots/%s-desktop.png' % theme_id)
-    browser.set_window_size(360, 740)  # Samsung Galaxy S8
-    browser.save_screenshot('shots/%s-mobile.png' % theme_id)
-    browser.set_window_size(740, 360)  # ditto, horizontal
-    browser.save_screenshot('shots/%s-mobile-horizontal.png' % theme_id)
+    shots = [
+        [1400, 1000, 'desktop'],
+        [360, 740, 'mobile'],  # Samsung Galaxy S8
+        [740, 360, 'mobile-horieontal'],  # ditto, horizonal
+    ]
+    if not args.reshot:
+        if all([os.path.exists(os.path.join(args.directory, '%s-%s.png' % (theme_id, x[2]))) for x in shots]):
+            sys.stderr.write(u' 📷 \n')
+            continue
+    if get_selected_theme() != theme_id:
+        set_theme(theme_id)
+        notify_agents(None)
+        for i in range(args.timeout):
+            sys.stderr.write('.')
+            time.sleep(1)
+    if not args.no_memcached:
+        os.system('sudo service memcached restart')
+    browser.get(args.url)
+    for shot in shots:
+        browser.set_window_size(shot[0], shot[1])
+        browser.save_screenshot(os.path.join(args.directory, '%s-%s.png' % (theme_id, shot[2])))
     sys.stderr.write(u' 📸 \n')
 
 browser.close()