Merge pull request #63 from edanalytics/feature/cli_yaml_overrides

adds a `--set` flag to the CLI to enable overriding config values in `lightbeam.yml`

Author: tomreitz

lightbeam/__main__.py

@@ -108,6 +108,11 @@ def main(argv=None):
         type=str,
         help='produces a JSON output file with structured information about run results'
     )
+    parser.add_argument("--set",
+        type=str,
+        nargs="*",
+        help='overrides a setting in the config YAML; example: --set fetch.page_size 1000'
+    )
 
     defaults = { "selector":"*", "params": "", "older_than": "", "newer_than": "", "resend_status_codes": "", "results_file": "" }
     parser.set_defaults(**defaults)
@@ -143,6 +148,12 @@ def main(argv=None):
 
     if not args.config_file:
         logger.error("config file not specified with `-c` flag, and no default {" + ", ".join(DEFAULT_CONFIG_FILES) + "} found")
+    
+    if args.set and len(args.set)%2 != 0: # odd number of overrides
+        logger.error("overrides specified with --set must be followed by an even number of strings (key value key value ...)")
+    overrides = None
+    if args.set:
+        overrides = dict(zip(args.set[::2], args.set[1::2]))
 
     lb = Lightbeam(
         config_file=args.config_file,
@@ -158,7 +169,8 @@ def main(argv=None):
         older_than=args.older_than,
         newer_than=args.newer_than,
         resend_status_codes=args.resend_status_codes,
-        results_file=args.results_file
+        results_file=args.results_file,
+        overrides=overrides,
         )
     try:
         logger.info("starting...")

lightbeam/lightbeam.py

@@ -55,7 +55,7 @@ class Lightbeam:
     MAX_STATUS_REASONS_TO_DISPLAY = 10
     DATA_FILE_EXTENSIONS = ['json', 'jsonl', 'ndjson']
 
-    def __init__(self, config_file, logger=None, selector="*", exclude="", keep_keys="*", drop_keys="", query="{}", params="", wipe=False, force=False, older_than="", newer_than="", resend_status_codes="", results_file=""):
+    def __init__(self, config_file, logger=None, selector="*", exclude="", keep_keys="*", drop_keys="", query="{}", params="", wipe=False, force=False, older_than="", newer_than="", resend_status_codes="", results_file="", overrides={}):
         self.config_file = config_file
         self.logger = logger
         self.errors = 0
@@ -82,12 +82,16 @@ def __init__(self, config_file, logger=None, selector="*", exclude="", keep_keys
         self.token_version = 0        
         self.results_file = os.path.abspath(results_file) if results_file else None
         self.start_timestamp = datetime.now()
+        self.overrides = overrides
 
         # load params and/or env vars for config YAML interpolation
         self.params = json.loads(params) if params else {}
         user_config = self.load_config_file()
-        
         self.config = util.merge_dicts(user_config, self.config_defaults)
+        
+        # inject overrides into config
+        if self.overrides: self.inject_cli_overrides()
+
         if "state_dir" in self.config:
             self.track_state = True
             self.config["state_dir"] = os.path.expanduser(self.config["state_dir"])
@@ -129,6 +133,39 @@ def __init__(self, config_file, logger=None, selector="*", exclude="", keep_keys
             "resources": {}
         }
 
+    def inject_cli_overrides(self):
+        # parse self.overrides into configs:
+        for key, value in self.overrides.items():
+            self.config = Lightbeam.set_path(self.config, key, value)
+
+    @staticmethod
+    def set_path(my_dict, path, value):
+        path_pieces = path.split(".")
+        current = my_dict
+        for path_piece in path_pieces[:-1]:
+            if path_piece not in current.keys():
+                current[path_piece] = {}
+            current = current[path_piece]
+        current[path_pieces[-1]] = Lightbeam.autocast(value)
+        return my_dict
+    
+    @staticmethod
+    def autocast(value):
+        if value.lower() in ['true', 'yes', 'on', 't', 'y']:
+            return True
+        elif value.lower() in ['false', 'no', 'off', 'f', 'n']:
+            return False
+        elif '.' in value:
+            try:
+                return float(value)
+            except ValueError:
+                return value
+        else:
+            try:
+                return int(value)
+            except ValueError:
+                return value
+    
     # this is intended to be called before any CRITICAL errors;
     # any cleanup tasks should go here:
     def shutdown(self, method):