fix #578 : allowed to save and load all Axis and Group objects of a session in/from HDF, CSV and EXCEL files

Author: alixdamman

doc/source/changes/version_0_29.rst.inc

@@ -43,6 +43,9 @@ Miscellaneous improvements
     >>> s2
     Session(arr1, arr2, a, b, a01)
 
+  Note: all axes (groups) of a session are stored in the same CSV file/Excel sheet/HDF group
+  named __axes__ (__groups__).
+
 * added examples for `read_excel` and `read_hdf` functions (closes :issue:`617`).
 
 

larray/core/axis.py

@@ -12,7 +12,7 @@
 from larray.core.abstractbases import ABCAxis, ABCAxisReference, ABCLArray
 from larray.core.expr import ExprNode
 from larray.core.group import (Group, LGroup, IGroup, IGroupMaker, _to_tick, _to_ticks, _to_key, _seq_summary,
-                               _contain_group_ticks, _seq_group_to_name, _translate_group_key_hdf)
+                               _seq_group_to_name, _translate_group_key_hdf)
 from larray.util.oset import *
 from larray.util.misc import (basestring, PY2, unicode, long, duplicates, array_lookup2, ReprString, index_by_id,
                               renamed_to, common_type, LHDFStore)

larray/core/group.py

@@ -1376,6 +1376,14 @@ def __getattr__(self, key):
         else:
             return getattr(self.eval(), key)
 
+    # needed to make *un*pickling work (because otherwise, __getattr__ is called before .key exists, which leads to
+    # an infinite recursion)
+    def __getstate__(self):
+        return self.__dict__
+
+    def __setstate__(self, d):
+        self.__dict__ = d
+
     def __hash__(self):
         # to_tick & to_key are partially opposite operations but this standardize on a single notation so that they can
         # all target each other. eg, this removes spaces in "list strings", instead of hashing them directly

larray/core/session.py

@@ -12,7 +12,11 @@
 from larray.core.axis import Axis
 from larray.core.array import LArray, get_axes, ndtest, zeros, zeros_like, sequence, aslarray
 from larray.util.misc import float_error_handler_factory, is_interactive_interpreter, renamed_to, inverseop
-from larray.inout.session import check_pattern, handler_classes, ext_default_engine
+from larray.inout.session import ext_default_engine, get_file_handler
+
+
+def check_pattern(k, pattern):
+    return k.startswith(pattern)
 
 
 # XXX: inherit from OrderedDict or LArray?
@@ -253,8 +257,11 @@ def load(self, fname, names=None, engine='auto', display=False, **kwargs):
             _, ext = os.path.splitext(fname)
             ext = ext.strip('.') if '.' in ext else 'csv'
             engine = ext_default_engine[ext]
-        handler_cls = handler_classes[engine]
-        handler = handler_cls(fname)
+        handler_cls = get_file_handler(engine)
+        if engine == 'pandas_csv' and 'sep' in kwargs:
+            handler = handler_cls(fname, kwargs['sep'])
+        else:
+            handler = handler_cls(fname)
         objects = handler.read_items(names, display=display, **kwargs)
         for k, v in objects.items():
             self[k] = v
@@ -275,11 +282,15 @@ def save(self, fname, names=None, engine='auto', overwrite=True, display=False,
         engine : {'auto', 'pandas_csv', 'pandas_hdf', 'pandas_excel', 'xlwings_excel', 'pickle'}, optional
             Dump using `engine`. Defaults to 'auto' (use default engine for the format guessed from the file extension).
         overwrite: bool, optional
-            Whether or not to overwrite an existing file, if any. Ignored for CSV files. If False, file is updated.
-            Defaults to True.
+            Whether or not to overwrite an existing file, if any. Ignored for CSV files and 'pandas_excel' engine.
+            If False, file is updated. Defaults to True.
         display : bool, optional
             Whether or not to display which file is being worked on. Defaults to False.
 
+        Notes
+        -----
+        See Notes section from :py:meth:`~Session.to_csv` and :py:meth:`~Session.to_excel`.
+
         Examples
         --------
         >>> # axes
@@ -309,12 +320,12 @@ def save(self, fname, names=None, engine='auto', overwrite=True, display=False,
             _, ext = os.path.splitext(fname)
             ext = ext.strip('.') if '.' in ext else 'csv'
             engine = ext_default_engine[ext]
-        handler_cls = handler_classes[engine]
-        handler = handler_cls(fname, overwrite)
-        if engine != 'pandas_hdf':
-            items = self.filter(kind=LArray).items()
+        handler_cls = get_file_handler(engine)
+        if engine == 'pandas_csv' and 'sep' in kwargs:
+            handler = handler_cls(fname, overwrite, kwargs['sep'])
         else:
-            items = self.items()
+            handler = handler_cls(fname, overwrite)
+        items = self.items()
         if names is not None:
             names_set = set(names)
             items = [(k, v) for k, v in items if k in names_set]
@@ -387,7 +398,7 @@ def to_globals(self, names=None, depth=0, warn=True, inplace=False):
 
     def to_pickle(self, fname, names=None, overwrite=True, display=False, **kwargs):
         """
-        Dumps all array objects from the current session to a file using pickle.
+        Dumps LArray, Axis and Group objects from the current session to a file using pickle.
 
         WARNING: never load a pickle file (.pkl or .pickle) from an untrusted source, as it can lead to arbitrary code
         execution.
@@ -397,7 +408,8 @@ def to_pickle(self, fname, names=None, overwrite=True, display=False, **kwargs):
         fname : str
             Path for the dump.
         names : list of str or None, optional
-            List of names of objects to dump. Defaults to all objects present in the Session.
+            Names of LArray/Axis/Group objects to dump.
+            Defaults to all objects present in the Session.
         overwrite: bool, optional
             Whether or not to overwrite an existing file, if any.
             If False, file is updated. Defaults to True.
@@ -406,16 +418,21 @@ def to_pickle(self, fname, names=None, overwrite=True, display=False, **kwargs):
 
         Examples
         --------
-        >>> arr1, arr2, arr3 = ndtest((2, 2)), ndtest(4), ndtest((3, 2))   # doctest: +SKIP
-        >>> s = Session([('arr1', arr1), ('arr2', arr2), ('arr3', arr3)])  # doctest: +SKIP
+        >>> # axes
+        >>> a, b = Axis("a=a0..a2"), Axis("b=b0..b2")    # doctest: +SKIP
+        >>> # groups
+        >>> a01 = a['a0,a1'] >> 'a01'                    # doctest: +SKIP
+        >>> # arrays
+        >>> arr1, arr2 = ndtest((a, b)), ndtest(a)       # doctest: +SKIP
+        >>> s = Session([('a', a), ('b', b), ('a01', a01), ('arr1', arr1), ('arr2', arr2)])  # doctest: +SKIP
 
         Save all arrays
 
         >>> s.to_pickle('output.pkl')  # doctest: +SKIP
 
-        Save only some arrays
+        Save only some objects
 
-        >>> s.to_pickle('output.pkl', ['arr1', 'arr3'])  # doctest: +SKIP
+        >>> s.to_pickle('output.pkl', ['a', 'b', 'arr1'])  # doctest: +SKIP
         """
         self.save(fname, names, ext_default_engine['pkl'], overwrite, display, **kwargs)
 
@@ -462,61 +479,85 @@ def to_hdf(self, fname, names=None, overwrite=True, display=False, **kwargs):
 
     def to_excel(self, fname, names=None, overwrite=True, display=False, **kwargs):
         """
-        Dumps all array objects from the current session to an Excel file.
+        Dumps LArray, Axis and Group objects from the current session to an Excel file.
 
         Parameters
         ----------
         fname : str
-            Path for the dump.
+            Path of the file for the dump.
         names : list of str or None, optional
-            List of names of objects to dump. Defaults to all objects present in the Session.
+            Names of LArray/Axis/Group objects to dump.
+            Defaults to all objects present in the Session.
         overwrite: bool, optional
             Whether or not to overwrite an existing file, if any. If False, file is updated. Defaults to True.
         display : bool, optional
             Whether or not to display which file is being worked on. Defaults to False.
 
+        Notes
+        -----
+        - each array is saved in a separate sheet
+        - all Axis objects are saved together in the same sheet named __axes__
+        - all Group objects are saved together in the same sheet named __groups__
+
         Examples
         --------
-        >>> arr1, arr2, arr3 = ndtest((2, 2)), ndtest(4), ndtest((3, 2))   # doctest: +SKIP
-        >>> s = Session([('arr1', arr1), ('arr2', arr2), ('arr3', arr3)])  # doctest: +SKIP
+        >>> # axes
+        >>> a, b = Axis("a=a0..a2"), Axis("b=b0..b2")    # doctest: +SKIP
+        >>> # groups
+        >>> a01 = a['a0,a1'] >> 'a01'                    # doctest: +SKIP
+        >>> # arrays
+        >>> arr1, arr2 = ndtest((a, b)), ndtest(a)       # doctest: +SKIP
+        >>> s = Session([('a', a), ('b', b), ('a01', a01), ('arr1', arr1), ('arr2', arr2)])  # doctest: +SKIP
 
         Save all arrays
 
         >>> s.to_excel('output.xlsx')  # doctest: +SKIP
 
-        Save only some arrays
+        Save only some objects
 
-        >>> s.to_excel('output.xlsx', ['arr1', 'arr3'])  # doctest: +SKIP
+        >>> s.to_excel('output.xlsx', ['a', 'b', 'arr1'])  # doctest: +SKIP
         """
         self.save(fname, names, ext_default_engine['xlsx'], overwrite, display, **kwargs)
 
     dump_excel = renamed_to(to_excel, 'dump_excel')
 
     def to_csv(self, fname, names=None, display=False, **kwargs):
         """
-        Dumps all array objects from the current session to CSV files.
+        Dumps LArray, Axis and Group objects from the current session to CSV files.
 
         Parameters
         ----------
         fname : str
             Path for the directory that will contain CSV files.
         names : list of str or None, optional
-            List of names of objects to dump. Defaults to all objects present in the Session.
+            Names of LArray/Axis/Group objects to dump.
+            Defaults to all objects present in the Session.
         display : bool, optional
             Whether or not to display which file is being worked on. Defaults to False.
 
+        Notes
+        -----
+        - each array is saved in a separate file
+        - all Axis objects are saved together in the same CSV file named __axes__.csv
+        - all Group objects are saved together in the same CSV file named __groups__.csv
+
         Examples
         --------
-        >>> arr1, arr2, arr3 = ndtest((2, 2)), ndtest(4), ndtest((3, 2))   # doctest: +SKIP
-        >>> s = Session([('arr1', arr1), ('arr2', arr2), ('arr3', arr3)])  # doctest: +SKIP
+        >>> # axes
+        >>> a, b = Axis("a=a0..a2"), Axis("b=b0..b2")    # doctest: +SKIP
+        >>> # groups
+        >>> a01 = a['a0,a1'] >> 'a01'                    # doctest: +SKIP
+        >>> # arrays
+        >>> arr1, arr2 = ndtest((a, b)), ndtest(a)       # doctest: +SKIP
+        >>> s = Session([('a', a), ('b', b), ('a01', a01), ('arr1', arr1), ('arr2', arr2)])  # doctest: +SKIP
 
         Save all arrays
 
         >>> s.to_csv('./Output')  # doctest: +SKIP
 
         Save only some arrays
 
-        >>> s.to_csv('./Output', ['arr1', 'arr3'])  # doctest: +SKIP
+        >>> s.to_csv('./Output', ['a', 'b', 'arr1'])  # doctest: +SKIP
         """
         self.save(fname, names, ext_default_engine['csv'], display=display, **kwargs)
 

larray/inout/__init__.py

@@ -6,4 +6,5 @@
 from larray.inout.excel import *
 from larray.inout.hdf import *
 from larray.inout.sas import *
+from larray.inout.pickle import *
 from larray.inout.xw_excel import *

larray/inout/common.py

@@ -3,8 +3,6 @@
 import os
 from collections import OrderedDict
 
-from larray.core.axis import Axis
-from larray.core.group import Group
 from larray.core.array import LArray
 
 
@@ -23,9 +21,6 @@ def _get_index_col(nb_axes=None, index_col=None, wide=True):
     return index_col
 
 
-_allowed_types = (LArray, Axis, Group)
-
-
 class FileHandler(object):
     """
     Abstract class defining the methods for "file handler" subclasses.
@@ -51,16 +46,18 @@ def _open_for_read(self):
     def _open_for_write(self):
         raise NotImplementedError()
 
-    def list(self):
+    def list_items(self):
         """
-        Returns the list of objects' names.
+        Return list containing pairs (name, type) for all stored objects
         """
         raise NotImplementedError()
 
-    def _read_item(self, key, *args, **kwargs):
+    def _read_item(self, key, type, *args, **kwargs):
+        """Read item"""
         raise NotImplementedError()
 
-    def _dump(self, key, value, *args, **kwargs):
+    def _dump_item(self, key, value, *args, **kwargs):
+        """Dump item. Raises an TypeError if type not taken into account by the FileHandler subclass."""
         raise NotImplementedError()
 
     def save(self):
@@ -111,14 +108,15 @@ def read_items(self, keys, *args, **kwargs):
         display = kwargs.pop('display', False)
         ignore_exceptions = kwargs.pop('ignore_exceptions', False)
         self._open_for_read()
+        key_types = self.list_items()
+        if keys is not None:
+            key_types = [(key, type) for key, type in key_types if key in keys]
         res = OrderedDict()
-        if keys is None:
-            keys = self.list()
-        for key in keys:
+        for key, type in key_types:
             if display:
-                print("loading", key, "...", end=' ')
+                print("loading", type, "object", key, "...", end=' ')
             try:
-                key, item = self._read_item(key, *args, **kwargs)
+                key, item = self._read_item(key, type, *args, **kwargs)
                 res[key] = item
             except Exception:
                 if not ignore_exceptions:
@@ -142,17 +140,20 @@ def dump_items(self, key_values, *args, **kwargs):
         display = kwargs.pop('display', False)
         self._get_original_file_name()
         self._open_for_write()
-        key_values = [(k, v) for k, v in key_values if isinstance(v, _allowed_types)]
         for key, value in key_values:
             if isinstance(value, LArray) and value.ndim == 0:
                 if display:
                     print('Cannot dump {}. Dumping 0D arrays is currently not supported.'.format(key))
                 continue
-            if display:
-                print("dumping", key, "...", end=' ')
-            self._dump(key, value, *args, **kwargs)
-            if display:
-                print("done")
+            try:
+                if display:
+                    print("dumping", key, "...", end=' ')
+                self._dump_item(key, value, *args, **kwargs)
+                if display:
+                    print("done")
+            except TypeError:
+                if display:
+                    print("Cannot dump {}. {} is not a supported type".format(key, type(value).__name__))
         self.save()
         self.close()
-        self._update_original_file()
+        self._update_original_file()