diff options
| -rw-r--r-- | crocoite/cli.py | 20 | ||||
| -rw-r--r-- | crocoite/task.py | 9 |
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, |