From c6cde93f4ee8afe0256b8b46e6f331375551921b Mon Sep 17 00:00:00 2001
From: Romain Dorgueil <romain@dorgueil.net>
Date: Sat, 30 Sep 2017 11:38:36 +0200
Subject: [PATCH] [cli] small refactoring in bonobo convert to use shortcuts as
 extensions if nothing else matched.

---
 bonobo/commands/convert.py | 42 ++++++++++++++++++++++++++------------
 1 file changed, 29 insertions(+), 13 deletions(-)

diff --git a/bonobo/commands/convert.py b/bonobo/commands/convert.py
index 90bd57b..3c26724 100644
--- a/bonobo/commands/convert.py
+++ b/bonobo/commands/convert.py
@@ -1,13 +1,15 @@
 import mimetypes
+import os
 
 import bonobo
 
 SHORTCUTS = {
-    'plain': 'text/plain',
-    'txt': 'text/plain',
-    'text': 'text/plain',
     'csv': 'text/csv',
     'json': 'application/json',
+    'pickle': 'pickle',
+    'plain': 'text/plain',
+    'text': 'text/plain',
+    'txt': 'text/plain',
 }
 
 REGISTRY = {
@@ -17,6 +19,9 @@ REGISTRY = {
     'text/plain': (bonobo.FileReader, bonobo.FileWriter),
 }
 
+READER = 'reader'
+WRITER = 'writer'
+
 
 def resolve_factory(name, filename, factory_type):
     """
@@ -30,32 +35,43 @@ def resolve_factory(name, filename, factory_type):
     if name in SHORTCUTS:
         name = SHORTCUTS[name]
 
-    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 name is None:
+        _, _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]
-    elif factory_type == 'writer':
+    elif factory_type == WRITER:
         return REGISTRY[name][1]
     else:
         raise ValueError('Invalid factory type.')
 
+
 def execute(input, output, reader=None, reader_options=None, writer=None, writer_options=None, options=None):
-    reader = resolve_factory(reader, input, 'reader')(input)
-    writer = resolve_factory(writer, output, 'writer')(output)
+    reader = resolve_factory(reader, input, READER)(input)
+    writer = resolve_factory(writer, output, WRITER)(output)
 
     graph = bonobo.Graph()
     graph.add_chain(reader, writer)
 
     return bonobo.run(graph, services={
-            'fs': bonobo.open_fs(),
-        })
+        'fs': bonobo.open_fs(),
+    })
+
 
 def register(parser):
     parser.add_argument('input')
     parser.add_argument('output')
-    parser.add_argument('--reader', '-r')
-    parser.add_argument('--writer', '-w')
+    parser.add_argument('--' + READER, '-r')
+    parser.add_argument('--' + WRITER, '-w')
     # parser.add_argument('--reader-option', '-ro', dest='reader_options')
     # parser.add_argument('--writer-option', '-wo', dest='writer_options')
     # parser.add_argument('--option', '-o', dest='options')