check good match of file type and io stream types

pyexcel · May 18, 2015 · a689be6 · a689be6
1 parent 3f759db
commit a689be6
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 6 deletions.
diff --git a/pyexcel_io/__init__.py b/pyexcel_io/__init__.py
@@ -23,7 +23,7 @@
 from .csvzipbook import CSVZipWriter, CSVZipBook
 from .sqlbook import SQLBookReader, SQLBookWriter
 from .djangobook import DjangoBookReader, DjangoBookWriter
-from ._compact import is_string, BytesIO, StringIO, isstream, OrderedDict, PY2
+from ._compact import is_string, BytesIO, StringIO, isstream, OrderedDict
 from .constants import (
     MESSAGE_LOADING_FORMATTER,
     MESSAGE_ERROR_03,
@@ -45,6 +45,10 @@
 )
 
 
+BINARY_STREAM_TYPES = [FILE_FORMAT_CSVZ, FILE_FORMAT_TSVZ,
+                       FILE_FORMAT_ODS, FILE_FORMAT_XLS,
+                       FILE_FORMAT_XLSX, FILE_FORMAT_XLSM]
+
 # A list of registered readers
 READERS = {
     FILE_FORMAT_CSV: CSVBook,
@@ -124,7 +128,10 @@ def load_data(filename,
             book_class = READERS[extension]
             if from_memory:
                 if isstream(filename):
-                    content = filename
+                    if validate_io(file_type, filename):
+                        content = filename
+                    else:
+                        raise IOError(MESSAGE_ERROR_03)
                 else:
                     io = get_io(file_type)
                     io.write(filename)
@@ -166,6 +173,8 @@ def get_writer(filename, file_type=None, **keywords):
             if isstream(filename):
                 extension = file_type
                 to_memory = True
+                if not validate_io(file_type, filename):
+                    raise IOError(MESSAGE_ERROR_03)
             else:
                 raise IOError(MESSAGE_ERROR_03)
         elif is_string(type(filename)):
@@ -187,12 +196,21 @@ def get_writer(filename, file_type=None, **keywords):
 def get_io(file_type):
     if file_type in [FILE_FORMAT_CSV, FILE_FORMAT_TSV]:
         return StringIO()
-    elif file_type in [FILE_FORMAT_CSVZ, FILE_FORMAT_TSVZ, FILE_FORMAT_ODS, FILE_FORMAT_XLS, FILE_FORMAT_XLSX, FILE_FORMAT_XLSM]:
+    elif file_type in BINARY_STREAM_TYPES:
         return BytesIO()
     else:
         return None
 
 
+def validate_io(file_type, io):
+    if file_type in [FILE_FORMAT_CSV, FILE_FORMAT_TSV]:
+        return isinstance(io, StringIO)
+    elif file_type in BINARY_STREAM_TYPES:
+        return isinstance(io, BytesIO)
+    else:
+        return False
+
+
 def store_data(afile, data, file_type=None, **keywords):
     """Non public function to store data to afile
 

diff --git a/tests/test_io.py b/tests/test_io.py
@@ -1,4 +1,4 @@
-from pyexcel_io import load_data, StringIO, get_writer, get_io
+from pyexcel_io import load_data, StringIO, get_writer, get_io, BytesIO
 from nose.tools import raises
 
 
@@ -24,21 +24,32 @@ def test_load_ods_data():
 
 @raises(NotImplementedError)
 def test_load_xls_data_from_memory():
-    io = StringIO()
+    io = BytesIO()
     load_data(io, file_type="xls")
 
+@raises(IOError)
+def test_load_xlsm_data_from_memory():
+    io = StringIO()
+    get_writer(io, file_type="xlsm")
 
+
 @raises(NotImplementedError)
 def test_write_xlsx_data():
     get_writer("test.xlsx")
 
 
-@raises(NotImplementedError)
+@raises(IOError)
 def test_writer_xlsm_data_from_memory():
     io = StringIO()
     get_writer(io, file_type="xlsm")
 
 
+@raises(NotImplementedError)
+def test_writer_xlsm_data_from_memory2():
+    io = BytesIO()
+    get_writer(io, file_type="xlsm")
+
+
 def test_get_io():
     io = get_io("hello")
     assert io == None