+class Config:
+ """
+ Everything in the config module used to be module-level functions and
+ variables but it made the code ugly and harder to maintain. Now, this
+ class does the heavy lifting. We still rely on some globals, though:
+
+ ARGS and GROUP to interface with argparse
+ PROGRAM_NAME stores argv[0] close to program invocation
+ ORIG_ARGV stores the original argv list close to program invocation
+ CONFIG and config: hold the (singleton) instance of this class.
+
+ """
+
+ def __init__(self):
+ # Has our parse() method been invoked yet?
+ self.config_parse_called = False
+
+ # A configuration dictionary that will contain parsed
+ # arguments. This is the data that is most interesting to our
+ # callers as it will hold the configuration result.
+ self.config: Dict[str, Any] = {}
+
+ # Defer logging messages until later when logging has been
+ # initialized.
+ self.saved_messages: List[str] = []
+
+ # A zookeeper client that is lazily created so as to not incur
+ # the latency of connecting to zookeeper for programs that are
+ # not reading or writing their config data into zookeeper.
+ self.zk: Optional[KazooClient] = None
+
+ # Per known zk file, what is the max version we have seen?
+ self.max_version: Dict[str, int] = {}
+
+ def __getitem__(self, key: str) -> Optional[Any]:
+ """If someone uses []'s on us, pass it onto self.config."""
+ return self.config.get(key, None)
+
+ def __setitem__(self, key: str, value: Any) -> None:
+ self.config[key] = value
+
+ def __contains__(self, key: str) -> bool:
+ return key in self.config
+
+ def get(self, key: str, default: Any = None) -> Optional[Any]:
+ return self.config.get(key, default)
+
+ @staticmethod
+ def add_commandline_args(title: str, description: str = "") -> argparse._ArgumentGroup:
+ """Create a new context for arguments and return a handle.
+
+ Args:
+ title: A title for your module's commandline arguments group.
+ description: A helpful description of your module.
+
+ Returns:
+ An argparse._ArgumentGroup to be populated by the caller.
+ """
+ return ARGS.add_argument_group(title, description)
+
+ @staticmethod
+ def overwrite_argparse_epilog(msg: str) -> None:
+ """Allows your code to override the default epilog created by
+ argparse.
+
+ Args:
+ msg: The epilog message to substitute for the default.
+ """
+ ARGS.epilog = msg
+
+ @staticmethod
+ def is_flag_already_in_argv(var: str) -> bool:
+ """Returns true if a particular flag is passed on the commandline
+ and false otherwise.
+
+ Args:
+ var: The flag to search for.
+ """
+ for _ in sys.argv:
+ if var in _:
+ return True
+ return False
+
+ @staticmethod
+ def print_usage() -> None:
+ """Prints the normal help usage message out."""
+ ARGS.print_help()
+
+ @staticmethod
+ def usage() -> str:
+ """
+ Returns:
+ program usage help text as a string.
+ """
+ return ARGS.format_usage()
+
+ @staticmethod
+ def _reorder_arg_action_groups_before_help(entry_module: Optional[str]):
+ """Internal. Used to reorder the arguments before dumping out a
+ generated help string such that the main program's arguments come
+ last.
+
+ """
+ reordered_action_groups = []
+ for grp in ARGS._action_groups:
+ if entry_module is not None and entry_module in grp.title: # type: ignore
+ reordered_action_groups.append(grp)
+ elif PROGRAM_NAME in GROUP.title: # type: ignore
+ reordered_action_groups.append(grp)
+ else:
+ reordered_action_groups.insert(0, grp)
+ return reordered_action_groups
+
+ @staticmethod
+ def _parse_arg_into_env(arg: str) -> Optional[Tuple[str, str, List[str]]]:
+ """Internal helper to parse commandling args into environment vars."""
+ arg = arg.strip()
+ if not arg.startswith('['):
+ return None
+ arg = arg.strip('[')
+ if not arg.endswith(']'):
+ return None
+ arg = arg.strip(']')
+
+ chunks = arg.split()
+ if len(chunks) > 1:
+ var = chunks[0]
+ else:
+ var = arg
+
+ # Environment vars the same as flag names without
+ # the initial -'s and in UPPERCASE.
+ env = var.upper()
+ while env[0] == '-':
+ env = env[1:]
+ return var, env, chunks
+
+ def _augment_sys_argv_from_environment_variables(self):
+ """Internal. Look at the system environment for variables that match
+ commandline arg names. This is done via some munging such that:
+
+ :code:`--argument_to_match`
+
+ ...is matched by:
+
+ :code:`ARGUMENT_TO_MATCH`
+
+ This allows users to set args via shell environment variables
+ in lieu of passing them on the cmdline.
+
+ """
+ usage_message = Config.usage()
+ optional = False
+ arg = ''
+
+ # Foreach valid optional commandline option (chunk) generate
+ # its analogous environment variable.
+ for chunk in usage_message.split():
+ if chunk[0] == '[':
+ optional = True
+ if optional:
+ arg += f'{chunk} '
+ if chunk[-1] == ']':
+ optional = False
+ _ = Config._parse_arg_into_env(arg)
+ if _:
+ var, env, chunks = _
+ if env in os.environ:
+ if not Config.is_flag_already_in_argv(var):
+ value = os.environ[env]
+ self.saved_messages.append(
+ f'Initialized from environment: {var} = {value}'
+ )
+ from string_utils import to_bool
+
+ if len(chunks) == 1 and to_bool(value):
+ sys.argv.append(var)
+ elif len(chunks) > 1:
+ sys.argv.append(var)
+ sys.argv.append(value)
+ arg = ''
+
+ def _process_dynamic_args(self, event: WatchedEvent):
+ assert self.zk
+ logger = logging.getLogger(__name__)
+ contents, meta = self.zk.get(event.path, watch=self._process_dynamic_args)
+ logger.debug('Update for %s at version=%d.', event.path, meta.version)
+ logger.debug(
+ 'Max known version for %s is %d.', event.path, self.max_version.get(event.path, 0)
+ )
+ if meta.version > self.max_version.get(event.path, 0):
+ self.max_version[event.path] = meta.version
+ contents = contents.decode()
+ temp_argv = []
+ for arg in contents.split():
+ if 'dynamic' in arg:
+ temp_argv.append(arg)
+ logger.info("Updating %s from zookeeper async config change.", arg)
+ if len(temp_argv) > 0:
+ old_argv = sys.argv
+ sys.argv = temp_argv
+ known, _ = ARGS.parse_known_args()
+ sys.argv = old_argv
+ self.config.update(vars(known))
+
+ def _augment_sys_argv_from_loadfile(self):
+ """Internal. Augment with arguments persisted in a saved file."""
+
+ loadfile = None
+ saw_other_args = False
+ grab_next_arg = False
+ for arg in sys.argv[1:]:
+ if 'config_loadfile' in arg:
+ pieces = arg.split('=')
+ if len(pieces) > 1:
+ loadfile = pieces[1]