fix #842 : included scalars when dumping or loading a Session object (hdf5 + pickle formats)

Author: alixdamman

doc/source/changes/version_0_33.rst.inc

@@ -49,7 +49,8 @@ New features
 Miscellaneous improvements
 ^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-* improved something.
+* scalar objects (i.e of type int, float, bool, string, date, time or datetime) belonging to a session
+  are now also saved and loaded when using the HDF5 or pickle format (closes :issue:`842`).
 
 
 Fixes

larray/core/session.py

@@ -1,11 +1,26 @@
 from __future__ import absolute_import, print_function
 
 import os
+from datetime import date, time, datetime
 from collections import OrderedDict
 
+from larray.util.compat import bytes, unicode
+from larray.core.axis import Axis
+from larray.core.group import Group
 from larray.core.array import Array
 
 
+# all formats
+_supported_larray_types = (Axis, Group, Array)
+
+# only for HDF5 and pickle formats
+# support list, tuple and dict?
+# replace unicode by str when Python 2.7 will no longer be supported
+_supported_scalars_types = (int, float, bool, bytes, unicode, date, time, datetime)
+_supported_types = _supported_larray_types + _supported_scalars_types
+_supported_typenames = {cls.__name__ for cls in _supported_types}
+
+
 def _get_index_col(nb_axes=None, index_col=None, wide=True):
     if not wide:
         if nb_axes is not None or index_col is not None:

larray/inout/common.py

@@ -3,6 +3,7 @@
 import warnings
 
 import numpy as np
+import pandas as pd
 from pandas import HDFStore
 
 from larray.core.array import Array
@@ -12,21 +13,27 @@
 from larray.core.metadata import Metadata
 from larray.util.misc import LHDFStore
 from larray.inout.session import register_file_handler
-from larray.inout.common import FileHandler
+from larray.inout.common import FileHandler, _supported_typenames, _supported_scalars_types
 from larray.inout.pandas import df_asarray
 from larray.example import get_example_filepath
 
 
+# for backward compatibility (larray < 0.29) but any object read from an hdf file should have
+# an attribute 'type'
+def _get_type_from_attrs(attrs):
+    return attrs.type if 'type' in attrs else 'Array'
+
+
 def read_hdf(filepath_or_buffer, key, fill_value=nan, na=nan, sort_rows=False, sort_columns=False,
              name=None, **kwargs):
-    r"""Reads an axis or group or array named key from a HDF5 file in filepath (path+name)
+    r"""Reads a scalar or an axis or group or array named key from a HDF5 file in filepath (path+name)
 
     Parameters
     ----------
     filepath_or_buffer : str or pandas.HDFStore
         Path and name where the HDF5 file is stored or a HDFStore object.
     key : str or Group
-        Name of the array.
+        Name of the scalar or axis or group or array.
     fill_value : scalar or Array, optional
         Value used to fill cells corresponding to label combinations which are not present in the input.
         Defaults to NaN.
@@ -70,11 +77,14 @@ def read_hdf(filepath_or_buffer, key, fill_value=nan, na=nan, sort_rows=False, s
     key = _translate_group_key_hdf(key)
     res = None
     with LHDFStore(filepath_or_buffer) as store:
-        pd_obj = store.get(key)
+        try:
+            pd_obj = store.get(key)
+        except KeyError:
+            filepath = filepath_or_buffer if isinstance(filepath_or_buffer, HDFStore) else store.filename
+            raise KeyError('No item with name {} has been found in file {}'.format(key, filepath))
         attrs = store.get_storer(key).attrs
         writer = attrs.writer if 'writer' in attrs else None
-        # for backward compatibility but any object read from an hdf file should have an attribute 'type'
-        _type = attrs.type if 'type' in attrs else 'Array'
+        _type = _get_type_from_attrs(attrs)
         _meta = attrs.metadata if 'metadata' in attrs else None
         if _type == 'Array':
             # cartesian product is not necessary if the array was written by LArray
@@ -110,6 +120,10 @@ def read_hdf(filepath_or_buffer, key, fill_value=nan, na=nan, sort_rows=False, s
                 key = np.char.decode(key, 'utf-8')
             axis = read_hdf(filepath_or_buffer, attrs['axis_key'])
             res = LGroup(key=key, name=name, axis=axis)
+        elif _type in _supported_typenames:
+            res = pd_obj.values
+            assert len(res) == 1
+            res = res[0]
     return res
 
 
@@ -126,36 +140,37 @@ def _open_for_write(self):
 
     def list_items(self):
         keys = [key.strip('/') for key in self.handle.keys()]
+        items = [(key, _get_type_from_attrs(self.handle.get_storer(key).attrs)) for key in keys if '/' not in key]
+        # ---- for backward compatibility (LArray < 0.33) ----
         # axes
-        items = [(key.split('/')[-1], 'Axis') for key in keys if '__axes__' in key]
+        items += [(key.split('/')[-1], 'Axis_Backward_Comp') for key in keys if '__axes__' in key]
         # groups
-        items += [(key.split('/')[-1], 'Group') for key in keys if '__groups__' in key]
-        # arrays
-        items += [(key, 'Array') for key in keys if '/' not in key]
+        items += [(key.split('/')[-1], 'Group_Backward_Comp') for key in keys if '__groups__' in key]
         return items
 
-    def _read_item(self, key, type, *args, **kwargs):
-        if type == 'Array':
+    def _read_item(self, key, typename, *args, **kwargs):
+        if typename in _supported_typenames:
             hdf_key = '/' + key
-        elif type == 'Axis':
+        # ---- for backward compatibility (LArray < 0.33) ----
+        elif typename == 'Axis_Backward_Comp':
             hdf_key = '__axes__/' + key
-        elif type == 'Group':
+        elif typename == 'Group_Backward_Comp':
             hdf_key = '__groups__/' + key
         else:
             raise TypeError()
         return read_hdf(self.handle, hdf_key, *args, **kwargs)
 
     def _dump_item(self, key, value, *args, **kwargs):
-        if isinstance(value, Array):
-            hdf_key = '/' + key
-            value.to_hdf(self.handle, hdf_key, *args, **kwargs)
-        elif isinstance(value, Axis):
-            hdf_key = '__axes__/' + key
+        hdf_key = '/' + key
+        if isinstance(value, (Array, Axis)):
             value.to_hdf(self.handle, hdf_key, *args, **kwargs)
         elif isinstance(value, Group):
-            hdf_key = '__groups__/' + key
-            hdf_axis_key = '__axes__/' + value.axis.name
+            hdf_axis_key = '/' + value.axis.name
             value.to_hdf(self.handle, hdf_key, hdf_axis_key, *args, **kwargs)
+        elif isinstance(value, _supported_scalars_types):
+            s = pd.Series(data=value)
+            self.handle.put(hdf_key, s)
+            self.handle.get_storer(hdf_key).attrs.type = type(value).__name__
         else:
             raise TypeError()
 

larray/inout/hdf.py

@@ -9,7 +9,7 @@
 from larray.core.metadata import Metadata
 from larray.util.compat import pickle
 from larray.inout.session import register_file_handler
-from larray.inout.common import FileHandler
+from larray.inout.common import FileHandler, _supported_types, _supported_typenames, _supported_scalars_types
 
 
 @register_file_handler('pickle', ['pkl', 'pickle'])
@@ -25,22 +25,25 @@ def _open_for_write(self):
             self.data = OrderedDict()
 
     def list_items(self):
+        # scalar
+        items = [(key, type(value).__name__) for key, value in self.data.items()
+                 if isinstance(value, _supported_scalars_types)]
         # axes
-        items = [(key, 'Axis') for key, value in self.data.items() if isinstance(value, Axis)]
+        items += [(key, 'Axis') for key, value in self.data.items() if isinstance(value, Axis)]
         # groups
         items += [(key, 'Group') for key, value in self.data.items() if isinstance(value, Group)]
         # arrays
         items += [(key, 'Array') for key, value in self.data.items() if isinstance(value, Array)]
         return items
 
-    def _read_item(self, key, type, *args, **kwargs):
-        if type in {'Array', 'Axis', 'Group'}:
+    def _read_item(self, key, typename, *args, **kwargs):
+        if typename in _supported_typenames:
             return self.data[key]
         else:
             raise TypeError()
 
     def _dump_item(self, key, value, *args, **kwargs):
-        if isinstance(value, (Array, Axis, Group)):
+        if isinstance(value, _supported_types):
             self.data[key] = value
         else:
             raise TypeError()

larray/inout/pickle.py

@@ -2,13 +2,15 @@
 
 import os
 import shutil
+from datetime import date, time, datetime
 
 import numpy as np
 import pandas as pd
 import pytest
 
 from larray.tests.common import (assert_array_nan_equal, inputpath, tmp_path, meta,
                                  needs_xlwings, needs_pytables, needs_xlrd)
+from larray.inout.common import _supported_scalars_types
 from larray import (Session, Axis, Array, Group, isnan, zeros_like, ndtest, ones_like, ones, full,
                     local_arrays, global_arrays, arrays)
 from larray.util.compat import pickle, PY2
@@ -178,7 +180,7 @@ def test_names(session):
 def _test_io(fpath, session, meta, engine):
     is_excel_or_csv = 'excel' in engine or 'csv' in engine
 
-    kind = Array if is_excel_or_csv else (Axis, Group, Array)
+    kind = Array if is_excel_or_csv else (Axis, Group, Array) + _supported_scalars_types
     session = session.filter(kind=kind)
 
     session.meta = meta
@@ -226,8 +228,21 @@ def _test_io(fpath, session, meta, engine):
         assert s.meta == meta
 
 
+def _add_scalars_to_session(s):
+    # 's' for scalar
+    s['s_int'] = 5
+    s['s_float'] = 5.5
+    s['s_bool'] = True
+    s['s_str'] = 'string'
+    s['s_date'] = date(2020, 1, 10)
+    s['s_time'] = time(11, 23, 54)
+    s['s_datetime'] = datetime(2020, 1, 10, 11, 23, 54)
+    return s
+
+
 @needs_pytables
 def test_h5_io(tmpdir, session, meta):
+    session = _add_scalars_to_session(session)
     fpath = tmp_path(tmpdir, 'test_session.h5')
     _test_io(fpath, session, meta, engine='pandas_hdf')
 
@@ -276,6 +291,7 @@ def test_csv_io(tmpdir, session, meta):
 
 
 def test_pickle_io(tmpdir, session, meta):
+    session = _add_scalars_to_session(session)
     fpath = tmp_path(tmpdir, 'test_session.pkl')
     _test_io(fpath, session, meta, engine='pickle')