Merge pull request #40 from fact-project/fix_io_kwargs

Fix wrong keyword arguments to io functions

.travis.yml

@@ -1,7 +1,7 @@
 language: python
 python:
-  - "3.4"
   - "3.5"
+  - "3.6"
 before_install:
   - export DISPLAY=:99.0
   - sh -e /etc/init.d/xvfb start
@@ -14,7 +14,7 @@ before_install:
   - sudo ln -s /run/shm /dev/shm
 
 install:
-  - conda install --yes python=$TRAVIS_PYTHON_VERSION atlas numpy scipy matplotlib pip pymongo flake8 pandas astropy pygments
+  - conda install --yes python=$TRAVIS_PYTHON_VERSION atlas numpy scipy matplotlib pip pymongo flake8 pandas astropy pygments pytables
   - pip install restructuredtext-lint
   - pip install .
 

fact/io.py

@@ -25,24 +25,25 @@
 native_byteorder = native_byteorder = {'little': '<', 'big': '>'}[sys.byteorder]
 
 
-def write_data(df, file_path, key='data', use_hp5y=False):
+def write_data(df, file_path, key='data', use_hp5y=False, **kwargs):
 
     name, extension = path.splitext(file_path)
 
     if extension in ['.hdf', '.hdf5', '.h5']:
         if use_hp5y is True:
-            to_h5py(file_path, df, key=key)
+            to_h5py(file_path, df, key=key, **kwargs)
         else:
-            df.to_hdf(file_path, key=key, format='table')
+            df.to_hdf(file_path, key=key, **kwargs)
 
     elif extension == '.json':
-        df.to_json(file_path)
+        df.to_json(file_path, **kwargs)
 
     elif extension in ('.jsonl', '.jsonline'):
-        df.to_json(file_path, lines=True, orient='records')
+        df.to_json(file_path, lines=True, orient='records', **kwargs)
 
     elif extension == '.csv':
-        df.to_csv(file_path, delimiter=',', index=False)
+        index = kwargs.pop('index', False)
+        df.to_csv(file_path, index=index, **kwargs)
 
     else:
         raise IOError(

setup.py

@@ -3,7 +3,7 @@
 
 setup(
     name='pyfact',
-    version='0.9.1',
+    version='0.9.2',
     description='A module containing useful methods for working with fact',
     url='http://github.com/fact-project/pyfact',
     author='Maximilian Noethe, Dominik Neise',
@@ -37,6 +37,7 @@
         'peewee',
         'h5py',
         'wrapt',
+        'tables',  # pytables in anaconda
     ],
     zip_safe=False,
 )

tests/test_io.py

@@ -2,6 +2,7 @@
 import tempfile
 import numpy as np
 import h5py
+import pytest
 
 
 def test_to_h5py():
@@ -89,3 +90,76 @@ def test_to_h5py_append():
 
         for col in df_written.columns:
             assert all(df_read[col] == df_written[col])
+
+
+def test_write_data_csv():
+    from fact.io import write_data
+
+    df = pd.DataFrame({
+        'x': np.random.normal(size=50),
+        'N': np.random.randint(0, 10, dtype='uint8')
+    })
+
+    with tempfile.NamedTemporaryFile(suffix='.csv') as f:
+        write_data(df, f.name)
+
+
+def test_write_data_json():
+    from fact.io import write_data
+
+    df = pd.DataFrame({
+        'x': np.random.normal(size=50),
+        'N': np.random.randint(0, 10, dtype='uint8')
+    })
+
+    with tempfile.NamedTemporaryFile(suffix='.json') as f:
+        write_data(df, f.name)
+
+
+def test_write_data_jsonlines():
+    from fact.io import write_data
+
+    df = pd.DataFrame({
+        'x': np.random.normal(size=50),
+        'N': np.random.randint(0, 10, dtype='uint8')
+    })
+
+    with tempfile.NamedTemporaryFile(suffix='.jsonl') as f:
+        write_data(df, f.name)
+
+
+def test_write_data_pandas_hdf():
+    from fact.io import write_data
+
+    df = pd.DataFrame({
+        'x': np.random.normal(size=50),
+        'N': np.random.randint(0, 10, dtype='uint8')
+    })
+
+    with tempfile.NamedTemporaryFile(suffix='.hdf5') as f:
+        write_data(df, f.name, use_h5py=False)
+
+
+def test_write_data_h5py():
+    from fact.io import write_data
+
+    df = pd.DataFrame({
+        'x': np.random.normal(size=50),
+        'N': np.random.randint(0, 10, dtype='uint8')
+    })
+
+    with tempfile.NamedTemporaryFile(suffix='.hdf5') as f:
+        write_data(df, f.name, use_h5py=True)
+
+
+def test_write_data_root():
+    from fact.io import write_data
+
+    df = pd.DataFrame({
+        'x': np.random.normal(size=50),
+        'N': np.random.randint(0, 10, dtype='uint8')
+    })
+
+    with pytest.raises(IOError):
+        with tempfile.NamedTemporaryFile(suffix='.root') as f:
+            write_data(df, f.name)