Refactoring the runner to go more towards standard python, also adds the ability to use bonobo argument parser from standard python execution.
This commit is contained in:
Romain Dorgueil
@ -1,91 +1,19 @@
|
||||
import argparse
|
||||
import codecs
|
||||
import os
|
||||
import os.path
|
||||
import runpy
|
||||
import sys
|
||||
from contextlib import contextmanager
|
||||
|
||||
from bonobo import settings, logging, get_argument_parser, patch_environ
|
||||
from bonobo.constants import DEFAULT_SERVICES_FILENAME, DEFAULT_SERVICES_ATTR
|
||||
from bonobo.util import get_name
|
||||
from bonobo import settings, logging
|
||||
from bonobo.commands.base import BaseCommand, BaseGraphCommand
|
||||
|
||||
logger = logging.get_logger()
|
||||
|
||||
|
||||
class BaseCommand:
|
||||
@property
|
||||
def logger(self):
|
||||
try:
|
||||
return self._logger
|
||||
except AttributeError:
|
||||
self._logger = logging.get_logger(get_name(self))
|
||||
return self._logger
|
||||
|
||||
def add_arguments(self, parser):
|
||||
"""
|
||||
Entry point for subclassed commands to add custom arguments.
|
||||
"""
|
||||
pass
|
||||
|
||||
def handle(self, *args, **options):
|
||||
"""
|
||||
The actual logic of the command. Subclasses must implement this method.
|
||||
"""
|
||||
raise NotImplementedError('Subclasses of BaseCommand must provide a handle() method')
|
||||
|
||||
|
||||
class BaseGraphCommand(BaseCommand):
|
||||
required = True
|
||||
|
||||
def add_arguments(self, parser):
|
||||
# target arguments (cannot provide both).
|
||||
source_group = parser.add_mutually_exclusive_group(required=self.required)
|
||||
source_group.add_argument('file', nargs='?', type=str)
|
||||
source_group.add_argument('-m', dest='mod', type=str)
|
||||
|
||||
# add arguments to enforce system environment.
|
||||
parser = get_argument_parser(parser)
|
||||
|
||||
return parser
|
||||
|
||||
def _run_path(self, file):
|
||||
return runpy.run_path(file, run_name='__main__')
|
||||
|
||||
def _run_module(self, mod):
|
||||
return runpy.run_module(mod, run_name='__main__')
|
||||
|
||||
def read(self, *, file, mod, args=None, **options):
|
||||
_graph, _options = None, None
|
||||
|
||||
def _record(graph, **options):
|
||||
nonlocal _graph, _options
|
||||
_graph, _options = graph, options
|
||||
|
||||
with _override_runner(_record), patch_environ(options):
|
||||
_argv = sys.argv
|
||||
try:
|
||||
if file:
|
||||
sys.argv = [file] + list(args) if args else [file]
|
||||
self._run_path(file)
|
||||
elif mod:
|
||||
sys.argv = [mod, *(args or ())]
|
||||
self._run_module(mod)
|
||||
else:
|
||||
raise RuntimeError('No target provided.')
|
||||
finally:
|
||||
sys.argv = _argv
|
||||
|
||||
if _graph is None:
|
||||
raise RuntimeError('Could not find graph.')
|
||||
|
||||
return _graph, _options
|
||||
|
||||
def handle(self, *args, **options):
|
||||
pass
|
||||
|
||||
|
||||
def entrypoint(args=None):
|
||||
"""
|
||||
Main callable for "bonobo" entrypoint.
|
||||
|
||||
Will load commands from "bonobo.commands" entrypoints, using stevedore.
|
||||
|
||||
"""
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--debug', '-D', action='store_true')
|
||||
|
||||
@ -113,8 +41,7 @@ def entrypoint(args=None):
|
||||
mgr = ExtensionManager(namespace='bonobo.commands')
|
||||
mgr.map(register_extension)
|
||||
|
||||
parsed_args, remaining = parser.parse_known_args(args)
|
||||
parsed_args = parsed_args.__dict__
|
||||
parsed_args = parser.parse_args(args).__dict__
|
||||
|
||||
if parsed_args.pop('debug', False):
|
||||
settings.DEBUG.set(True)
|
||||
@ -123,45 +50,6 @@ def entrypoint(args=None):
|
||||
|
||||
logger.debug('Command: ' + parsed_args['command'] + ' Arguments: ' + repr(parsed_args))
|
||||
|
||||
# Get command handler
|
||||
# Get command handler, execute, rince.
|
||||
command = commands[parsed_args.pop('command')]
|
||||
|
||||
if len(remaining):
|
||||
command(_remaining_args=remaining, **parsed_args)
|
||||
else:
|
||||
command(**parsed_args)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _override_runner(runner):
|
||||
import bonobo
|
||||
_get_argument_parser = bonobo.get_argument_parser
|
||||
_run = bonobo.run
|
||||
try:
|
||||
def get_argument_parser(parser=None):
|
||||
return parser or argparse.ArgumentParser()
|
||||
|
||||
bonobo.get_argument_parser = get_argument_parser
|
||||
bonobo.run = runner
|
||||
|
||||
yield runner
|
||||
finally:
|
||||
bonobo.get_argument_parser = _get_argument_parser
|
||||
bonobo.run = _run
|
||||
|
||||
|
||||
def set_env_var(e, override=False):
|
||||
__escape_decoder = codecs.getdecoder('unicode_escape')
|
||||
ename, evalue = e.split('=', 1)
|
||||
|
||||
def decode_escaped(escaped):
|
||||
return __escape_decoder(escaped)[0]
|
||||
|
||||
if len(evalue) > 0:
|
||||
if evalue[0] == evalue[len(evalue) - 1] in ['"', "'"]:
|
||||
evalue = decode_escaped(evalue[1:-1])
|
||||
|
||||
if override:
|
||||
os.environ[ename] = evalue
|
||||
else:
|
||||
os.environ.setdefault(ename, evalue)
|
||||
command(**parsed_args)
|
||||
|
||||
129
bonobo/commands/base.py
Normal file
129
bonobo/commands/base.py
Normal file
@ -0,0 +1,129 @@
|
||||
import argparse
|
||||
import runpy
|
||||
import sys
|
||||
from contextlib import contextmanager
|
||||
|
||||
import bonobo.util.environ
|
||||
from bonobo import logging
|
||||
from bonobo.util.environ import get_argument_parser, parse_args
|
||||
from bonobo.util import get_name
|
||||
|
||||
|
||||
class BaseCommand:
|
||||
"""
|
||||
Base class for CLI commands.
|
||||
|
||||
"""
|
||||
|
||||
@property
|
||||
def logger(self):
|
||||
try:
|
||||
return self._logger
|
||||
except AttributeError:
|
||||
self._logger = logging.get_logger(get_name(self))
|
||||
return self._logger
|
||||
|
||||
def add_arguments(self, parser):
|
||||
"""
|
||||
Entry point for subclassed commands to add custom arguments.
|
||||
"""
|
||||
pass
|
||||
|
||||
def handle(self, *args, **options):
|
||||
"""
|
||||
The actual logic of the command. Subclasses must implement this method.
|
||||
"""
|
||||
raise NotImplementedError('Subclasses of BaseCommand must provide a handle() method')
|
||||
|
||||
|
||||
class BaseGraphCommand(BaseCommand):
|
||||
"""
|
||||
Base class for CLI commands that depends on a graph definition, either from a file or from a module.
|
||||
|
||||
"""
|
||||
required = True
|
||||
handler = None
|
||||
|
||||
def add_arguments(self, parser):
|
||||
# target arguments (cannot provide both).
|
||||
source_group = parser.add_mutually_exclusive_group(required=self.required)
|
||||
source_group.add_argument('file', nargs='?', type=str)
|
||||
source_group.add_argument('-m', dest='mod', type=str)
|
||||
|
||||
# add arguments to enforce system environment.
|
||||
parser = get_argument_parser(parser)
|
||||
|
||||
return parser
|
||||
|
||||
def parse_options(self, **options):
|
||||
return options
|
||||
|
||||
def handle(self, file, mod, **options):
|
||||
options = self.parse_options(**options)
|
||||
with self.read(file, mod, **options) as (graph, graph_execution_options, options):
|
||||
return self.do_handle(graph, **graph_execution_options, **options)
|
||||
|
||||
def do_handle(self, graph, **options):
|
||||
if not self.handler:
|
||||
raise RuntimeError('{} has no handler defined.'.format(get_name(self)))
|
||||
return self.handler(graph, **options)
|
||||
|
||||
@contextmanager
|
||||
def read(self, file, mod, **options):
|
||||
_graph, _graph_execution_options = None, None
|
||||
|
||||
def _record(graph, **graph_execution_options):
|
||||
nonlocal _graph, _graph_execution_options
|
||||
_graph, _graph_execution_options = graph, graph_execution_options
|
||||
|
||||
with _override_runner(_record), parse_args(options) as options:
|
||||
_argv = sys.argv
|
||||
try:
|
||||
if file:
|
||||
sys.argv = [file]
|
||||
self._run_path(file)
|
||||
elif mod:
|
||||
sys.argv = [mod]
|
||||
self._run_module(mod)
|
||||
else:
|
||||
raise RuntimeError('No target provided.')
|
||||
finally:
|
||||
sys.argv = _argv
|
||||
|
||||
if _graph is None:
|
||||
raise RuntimeError('Could not find graph.')
|
||||
|
||||
yield _graph, _graph_execution_options, options
|
||||
|
||||
def _run_path(self, file):
|
||||
return runpy.run_path(file, run_name='__main__')
|
||||
|
||||
def _run_module(self, mod):
|
||||
return runpy.run_module(mod, run_name='__main__')
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _override_runner(runner):
|
||||
"""
|
||||
Context manager that monkey patches `bonobo.run` function with our current command logic.
|
||||
|
||||
:param runner: the callable that will handle the `run()` logic.
|
||||
"""
|
||||
import bonobo
|
||||
|
||||
_get_argument_parser = bonobo.util.environ.get_argument_parser
|
||||
_run = bonobo.run
|
||||
try:
|
||||
# Original get_argument_parser would create or update an argument parser with environment options, but here we
|
||||
# already had them parsed so let's patch with something that creates an empty one instead.
|
||||
def get_argument_parser(parser=None):
|
||||
return parser or argparse.ArgumentParser()
|
||||
|
||||
bonobo.util.environ.get_argument_parser = get_argument_parser
|
||||
bonobo.run = runner
|
||||
|
||||
yield runner
|
||||
finally:
|
||||
# Restore our saved values.
|
||||
bonobo.util.environ.get_argument_parser = _get_argument_parser
|
||||
bonobo.run = _run
|
||||
@ -48,8 +48,17 @@ class ConvertCommand(BaseCommand):
|
||||
help='Add a named option to the writer factory.',
|
||||
)
|
||||
|
||||
def handle(self, input_filename, output_filename, reader=None, reader_option=None, writer=None, writer_option=None,
|
||||
option=None, transformation=None):
|
||||
def handle(
|
||||
self,
|
||||
input_filename,
|
||||
output_filename,
|
||||
reader=None,
|
||||
reader_option=None,
|
||||
writer=None,
|
||||
writer_option=None,
|
||||
option=None,
|
||||
transformation=None
|
||||
):
|
||||
reader_factory = default_registry.get_reader_factory_for(input_filename, format=reader)
|
||||
reader_options = _resolve_options((option or []) + (reader_option or []))
|
||||
|
||||
|
||||
@ -1,21 +1,15 @@
|
||||
import bonobo
|
||||
from bonobo.commands import BaseGraphCommand
|
||||
|
||||
OUTPUT_GRAPH = 'graphviz'
|
||||
|
||||
|
||||
class InspectCommand(BaseGraphCommand):
|
||||
handler = staticmethod(bonobo.inspect)
|
||||
|
||||
def add_arguments(self, parser):
|
||||
super(InspectCommand, self).add_arguments(parser)
|
||||
parser.add_argument('--graph', '-g', dest='output', action='store_const', const=OUTPUT_GRAPH)
|
||||
|
||||
def handle(self, output=None, **options):
|
||||
if output is None:
|
||||
raise ValueError('Output type must be provided (try --graph/-g).')
|
||||
|
||||
graph, params = self.read(**options)
|
||||
|
||||
if output == OUTPUT_GRAPH:
|
||||
print(graph._repr_dot_())
|
||||
else:
|
||||
raise NotImplementedError('Output type not implemented.')
|
||||
parser.add_argument('--graph', '-g', dest='format', action='store_const', const='graph')
|
||||
|
||||
def parse_options(self, **options):
|
||||
if not options.get('format'):
|
||||
raise RuntimeError('You must provide a format (try --graph).')
|
||||
return options
|
||||
|
||||
@ -6,6 +6,7 @@ from bonobo.commands import BaseGraphCommand
|
||||
|
||||
class RunCommand(BaseGraphCommand):
|
||||
install = False
|
||||
handler = staticmethod(bonobo.run)
|
||||
|
||||
def add_arguments(self, parser):
|
||||
super(RunCommand, self).add_arguments(parser)
|
||||
@ -16,7 +17,15 @@ class RunCommand(BaseGraphCommand):
|
||||
|
||||
parser.add_argument('--install', '-I', action='store_true')
|
||||
|
||||
def parse_options(self, *, quiet=False, verbose=False, install=False, **options):
|
||||
from bonobo import settings
|
||||
settings.QUIET.set_if_true(quiet)
|
||||
settings.DEBUG.set_if_true(verbose)
|
||||
self.install = install
|
||||
return options
|
||||
|
||||
def _run_path(self, file):
|
||||
# add install logic
|
||||
if self.install:
|
||||
if os.path.isdir(file):
|
||||
requirements = os.path.join(file, 'requirements.txt')
|
||||
@ -27,24 +36,12 @@ class RunCommand(BaseGraphCommand):
|
||||
return super()._run_path(file)
|
||||
|
||||
def _run_module(self, mod):
|
||||
# install not implemented for a module, not sure it even make sense.
|
||||
if self.install:
|
||||
raise RuntimeError('--install behaviour when running a module is not defined.')
|
||||
|
||||
return super()._run_module(mod)
|
||||
|
||||
def handle(self, quiet=False, verbose=False, install=False, _remaining_args=None, **options):
|
||||
from bonobo import settings
|
||||
|
||||
settings.QUIET.set_if_true(quiet)
|
||||
settings.DEBUG.set_if_true(verbose)
|
||||
self.install = install
|
||||
|
||||
graph, params = self.read(args=_remaining_args, **options)
|
||||
|
||||
params['plugins'] = set(params.pop('plugins', ())).union(set(options.pop('plugins', ())))
|
||||
|
||||
return bonobo.run(graph, **params)
|
||||
|
||||
|
||||
def register_generic_run_arguments(parser, required=True):
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user