Source code for law.parser
# coding: utf-8
"""
Helpers to extract useful information from the luigi command line parser.
"""
__all__ = []
from collections import OrderedDict
from argparse import ArgumentParser
import luigi
from law.logger import get_logger
from law.util import multi_match
logger = get_logger(__name__)
# cached objects
_root_task_cls = None
_root_task = None
_full_parser = None
_root_task_parser = None
_global_cmdline_args = None
_global_cmdline_values = None
[docs]def root_task_cls(task=None):
"""
Returns the class of the task that was triggered on the command line. The returned class is
cached. When *task* is defined and no root task class was cached yet, this methods acts as a
setter.
"""
global _root_task_cls
if not _root_task_cls:
if task:
_root_task_cls = task.__class__ if isinstance(task, luigi.Task) else task
logger.debug("set root task class to externally passed instance")
else:
luigi_parser = luigi.cmdline_parser.CmdlineParser.get_instance()
if not luigi_parser:
return None
_root_task_cls = luigi_parser._get_task_cls()
logger.debug("built root task class using luigi argument parser")
return _root_task_cls
[docs]def root_task(task=None):
"""
Returns the instance of the task that was triggered on the command line. The returned instance
is cached. When *task* is defined and no root task was cached yet, this methods acts as a
setter.
"""
global _root_task
if not _root_task:
if task:
_root_task = task
logger.debug("set root task to externally passed instance")
else:
luigi_parser = luigi.cmdline_parser.CmdlineParser.get_instance()
if not luigi_parser:
return None
_root_task = luigi_parser.get_task_obj()
logger.debug("built root task instance using luigi argument parser")
return _root_task
[docs]def full_parser():
"""
Returns the full *ArgumentParser* used by the luigi ``CmdlineParser``. The returned instance is
cached.
"""
global _full_parser
if not _full_parser:
luigi_parser = luigi.cmdline_parser.CmdlineParser.get_instance()
if not luigi_parser:
return None
# build the full argument parser with luigi helpers
root_task = luigi_parser.known_args.root_task
_full_parser = luigi_parser._build_parser(root_task=root_task)
logger.debug("built full luigi argument parser")
return _full_parser
[docs]def root_task_parser():
"""
Returns a new *ArgumentParser* instance that only contains parameter actions of the root task.
The returned instance is cached.
"""
global _root_task_parser
if not _root_task_parser:
luigi_parser = luigi.cmdline_parser.CmdlineParser.get_instance()
if not luigi_parser:
return None
root_task = luigi_parser.known_args.root_task
# get all root task parameter destinations
root_dests = []
for task_name, _, param_name, _ in luigi.task_register.Register.get_all_params():
if task_name == root_task:
root_dests.append(param_name)
# create a new parser and add all root actions
_root_task_parser = ArgumentParser(add_help=False)
for action in list(full_parser()._actions):
if not action.option_strings or action.dest in root_dests:
_root_task_parser._add_action(action)
logger.debug("built luigi argument parser for root task {}".format(root_task))
return _root_task_parser
[docs]def global_cmdline_args(exclude=None):
"""
Returns a dictionary with keys and string values of command line arguments that do not belong to
the root task. For bool parameters, such as ``--local-scheduler``, ``"True"`` is assumed if they
are used as flags, i.e., without a parameter value. The returned dict is cached. *exclude* can
be a list of argument names (with or without the leading ``"--"``) to be removed. Example:
.. code-block:: python
global_cmdline_args()
# -> {"--local-scheduler": "True", "--workers": "4"}
global_cmdline_args(exclude=["workers"])
# -> {"--local-scheduler": "True"}
"""
global _global_cmdline_args
if not _global_cmdline_args:
luigi_parser = luigi.cmdline_parser.CmdlineParser.get_instance()
if not luigi_parser:
return None
_global_cmdline_args = OrderedDict()
args = list(root_task_parser().parse_known_args(luigi_parser.cmdline_args)[1])
# expand bool flags
while args:
arg = args.pop(0)
# the argument must start with "--"
if not arg.startswith("--"):
raise Exception("global argument must start with '--', found '{}'".format(arg))
# get the corresponding value which is either part of the argument itself in the format
# "--arg=value" or passed in the next argument which must not start with "--" (in this
# case it is interpreted as a boolean "True" value)
if "=" in arg:
arg, value = arg.split("=", 1)
elif args and not args[0].startswith("--"):
value = args.pop(0)
else:
value = "True"
_global_cmdline_args[arg] = value
args = _global_cmdline_args
if exclude:
args = OrderedDict(args)
# treat arguments in exclude as patterns
_exclude = ["--" + pattern.lstrip("-") for pattern in exclude]
for arg in list(args.keys()):
if multi_match(arg, _exclude, mode=any):
args.pop(arg)
return args
[docs]def global_cmdline_values():
"""
Returns a dictionary of global command line arguments (computed with
:py:func:`global_cmdline_args`) to their current values. The returnd dictionary is cached.
Example:
.. code-block:: python
global_cmdline_values()
# -> {"core_local_scheduler": True}
"""
global _global_cmdline_values
if not _global_cmdline_values:
luigi_parser = luigi.cmdline_parser.CmdlineParser.get_instance()
if not luigi_parser:
return None
# go through all actions of the full luigi parser and compare option strings
# with the global cmdline args
parser = full_parser()
global_args = global_cmdline_args()
_global_cmdline_values = {}
for action in parser._actions:
if any(arg in action.option_strings for arg in global_args):
_global_cmdline_values[action.dest] = getattr(luigi_parser.known_args, action.dest)
return _global_cmdline_values
def _reset():
"""
Resets all singletons defined by the parser functions above.
"""
global _root_task
global _full_parser
global _root_task_parser
global _global_cmdline_args
global _global_cmdline_values
_root_task = None
_full_parser = None
_root_task_parser = None
_global_cmdline_args = None
_global_cmdline_values = None
