[cli] small refactoring in bonobo convert to use shortcuts as extensions if nothing else matched.

This commit is contained in:
Romain Dorgueil
2017-09-30 11:38:36 +02:00
parent 34dc53aa19
commit c6cde93f4e

View File

@ -1,13 +1,15 @@
import mimetypes import mimetypes
import os
import bonobo import bonobo
SHORTCUTS = { SHORTCUTS = {
'plain': 'text/plain',
'txt': 'text/plain',
'text': 'text/plain',
'csv': 'text/csv', 'csv': 'text/csv',
'json': 'application/json', 'json': 'application/json',
'pickle': 'pickle',
'plain': 'text/plain',
'text': 'text/plain',
'txt': 'text/plain',
} }
REGISTRY = { REGISTRY = {
@ -17,6 +19,9 @@ REGISTRY = {
'text/plain': (bonobo.FileReader, bonobo.FileWriter), 'text/plain': (bonobo.FileReader, bonobo.FileWriter),
} }
READER = 'reader'
WRITER = 'writer'
def resolve_factory(name, filename, factory_type): def resolve_factory(name, filename, factory_type):
""" """
@ -30,32 +35,43 @@ def resolve_factory(name, filename, factory_type):
if name in SHORTCUTS: if name in SHORTCUTS:
name = SHORTCUTS[name] name = SHORTCUTS[name]
if not name in REGISTRY: if name is None:
raise RuntimeError('Could not resolve {factory_type} factory for {filename} ({name}). Try providing it explicitely using -{opt} <format>.'.format(name=name, filename=filename, factory_type=factory_type, opt=factory_type[0])) _, _ext = os.path.splitext(filename)
if _ext:
_ext = _ext[1:]
if _ext in SHORTCUTS:
name = SHORTCUTS[_ext]
if factory_type == 'reader': if not name in REGISTRY:
raise RuntimeError(
'Could not resolve {factory_type} factory for {filename} ({name}). Try providing it explicitely using -{opt} <format>.'.format(
name=name, filename=filename, factory_type=factory_type, opt=factory_type[0]))
if factory_type == READER:
return REGISTRY[name][0] return REGISTRY[name][0]
elif factory_type == 'writer': elif factory_type == WRITER:
return REGISTRY[name][1] return REGISTRY[name][1]
else: else:
raise ValueError('Invalid factory type.') raise ValueError('Invalid factory type.')
def execute(input, output, reader=None, reader_options=None, writer=None, writer_options=None, options=None): def execute(input, output, reader=None, reader_options=None, writer=None, writer_options=None, options=None):
reader = resolve_factory(reader, input, 'reader')(input) reader = resolve_factory(reader, input, READER)(input)
writer = resolve_factory(writer, output, 'writer')(output) writer = resolve_factory(writer, output, WRITER)(output)
graph = bonobo.Graph() graph = bonobo.Graph()
graph.add_chain(reader, writer) graph.add_chain(reader, writer)
return bonobo.run(graph, services={ return bonobo.run(graph, services={
'fs': bonobo.open_fs(), 'fs': bonobo.open_fs(),
}) })
def register(parser): def register(parser):
parser.add_argument('input') parser.add_argument('input')
parser.add_argument('output') parser.add_argument('output')
parser.add_argument('--reader', '-r') parser.add_argument('--' + READER, '-r')
parser.add_argument('--writer', '-w') parser.add_argument('--' + WRITER, '-w')
# parser.add_argument('--reader-option', '-ro', dest='reader_options') # parser.add_argument('--reader-option', '-ro', dest='reader_options')
# parser.add_argument('--writer-option', '-wo', dest='writer_options') # parser.add_argument('--writer-option', '-wo', dest='writer_options')
# parser.add_argument('--option', '-o', dest='options') # parser.add_argument('--option', '-o', dest='options')