Share recursive argument parser

2 files changed, 15 insertions, 14 deletions

diff --git a/crocoite/cli.py b/crocoite/cli.py
index efd30ad..f6454da 100644
--- a/crocoite/cli.py
+++ b/crocoite/cli.py
@@ -29,6 +29,16 @@ from .controller import RecursiveController, defaultSettings, \
         ControllerSettings, DepthLimit, PrefixLimit
 from .browser import NullService, ChromeService
 
+def parseRecursive (recursive, url):
+    if recursive is None:
+        return DepthLimit (0)
+    elif recursive.isdigit ():
+        return DepthLimit (int (recursive))
+    elif recursive == 'prefix':
+        return PrefixLimit (url)
+    else:
+        raise ValueError ('Unsupported')
+
 def main ():
     parser = argparse.ArgumentParser(description='Save website to WARC using Google Chrome.')
     parser.add_argument('--browser', help='DevTools URL', metavar='URL')
@@ -63,13 +73,9 @@ def main ():
     else:
         logging.basicConfig (level=logging.INFO)
 
-        if args.recursive is None:
-            recursionPolicy = DepthLimit (0)
-        elif args.recursive.isdigit ():
-            recursionPolicy = DepthLimit (int (args.recursive))
-        elif args.recursive == 'prefix':
-            recursionPolicy = PrefixLimit (args.url)
-        else:
+        try:
+            recursionPolicy = parseRecursive (args.recursive, args.url)
+        except ValueError:
             parser.error ('Invalid argument for --recursive')
         service = ChromeService ()
         if args.browser:
diff --git a/crocoite/task.py b/crocoite/task.py
index 52d3b26..e93cfde 100644
--- a/crocoite/task.py
+++ b/crocoite/task.py
@@ -41,6 +41,7 @@ from celery.utils.log import get_task_logger
 from .browser import ChromeService
 from .controller import SinglePageController, ControllerSettings, RecursiveController, defaultSettings, DepthLimit, PrefixLimit
 from . import behavior
+from .cli import parseRecursive
 
 app = Celery ('crocoite.distributed')
 app.config_from_object('celeryconfig')
@@ -109,13 +110,7 @@ class DistributedRecursiveController (RecursiveController):
 def controller (self, url, settings, enabledBehaviorNames, recursive, concurrency):
     """ Recursive controller """
 
-    if recursive is None:
-        recursionPolicy = DepthLimit (0)
-    elif recursive.isdigit ():
-        recursionPolicy = DepthLimit (int (recursive))
-    elif recursive == 'prefix':
-        recursionPolicy = PrefixLimit (url)
-
+    recursionPolicy = parseRecursive (recursive, url)
     enabledBehavior = list (filter (lambda x: x.name in enabledBehaviorNames, behavior.available))
     settings = ControllerSettings (**settings)
     controller = DistributedRecursiveController (url, None, behavior=enabledBehavior,