added returned type to FileHandler methods (issue #864)

Author: alixdamman

larray/inout/common.py

@@ -2,9 +2,12 @@
 from datetime import date, time, datetime
 from collections import OrderedDict
 
+from typing import List, Tuple
+
 from larray.core.axis import Axis
 from larray.core.group import Group
 from larray.core.array import Array
+from larray.core.metadata import Metadata
 
 
 # all formats
@@ -17,7 +20,7 @@
 _supported_typenames = {cls.__name__ for cls in _supported_types}
 
 
-def _get_index_col(nb_axes=None, index_col=None, wide=True):
+def _get_index_col(nb_axes=None, index_col=None, wide=True) -> List[int]:
     if not wide:
         if nb_axes is not None or index_col is not None:
             raise ValueError("`nb_axes` or `index_col` argument cannot be used when `wide` argument is False")
@@ -57,7 +60,8 @@ def _open_for_read(self):
     def _open_for_write(self):
         raise NotImplementedError()
 
-    def list_items(self):
+    # FIXME : return an ordinary dict instead (Python < 3.6 no longer supported)
+    def list_items(self) -> List[Tuple[str, str]]:
         r"""
         Return list containing pairs (name, type) for all stored objects
         """
@@ -67,7 +71,7 @@ def _read_item(self, key, type, *args, **kwargs):
         r"""Read item"""
         raise NotImplementedError()
 
-    def _read_metadata(self):
+    def _read_metadata(self) -> Metadata:
         r"""Read metadata"""
         raise NotImplementedError()
 
@@ -102,7 +106,8 @@ def _update_original_file(self):
             os.remove(self.original_file_name)
             os.rename(self.fname, self.original_file_name)
 
-    def read(self, keys, *args, **kwargs):
+    # FIXME : replace OrderedDict by ordinary dict (Python < 3.6 no longer supported)
+    def read(self, keys, *args, **kwargs) -> Tuple[Metadata, OrderedDict]:
         r"""
         Reads file content (HDF, Excel, CSV, ...) and returns a dictionary containing loaded objects.
 

larray/inout/csv.py

@@ -6,6 +6,8 @@
 import pandas as pd
 import numpy as np
 
+from typing import List, Tuple
+
 from larray.core.array import Array, asarray
 from larray.core.constants import nan
 from larray.core.metadata import Metadata
@@ -273,7 +275,7 @@ def __init__(self, fname, overwrite_file=False, sep=','):
     def _get_original_file_name(self):
         pass
 
-    def _to_filepath(self, key):
+    def _to_filepath(self, key: str) -> str:
         if self.directory is not None:
             return os.path.join(self.directory, f'{key}.csv')
         else:
@@ -291,7 +293,7 @@ def _open_for_write(self):
                 if not os.path.isdir(self.directory):
                     raise ValueError(f"Path {self.directory} must represent a directory")
 
-    def list_items(self):
+    def list_items(self) -> List[Tuple[str, str]]:
         fnames = glob(self.pattern) if self.pattern is not None else []
         # drop directory
         fnames = [os.path.basename(fname) for fname in fnames]
@@ -300,7 +302,7 @@ def list_items(self):
         fnames = sorted([os.path.splitext(fname)[0] for fname in fnames])
         return [(name, 'Array') for name in fnames if name != '__metadata__']
 
-    def _read_item(self, key, type, *args, **kwargs):
+    def _read_item(self, key, type, *args, **kwargs) -> Array:
         if type == 'Array':
             return read_csv(self._to_filepath(key), *args, **kwargs)
         else:
@@ -312,7 +314,7 @@ def _dump_item(self, key, value, *args, **kwargs):
         else:
             raise TypeError()
 
-    def _read_metadata(self):
+    def _read_metadata(self) -> Metadata:
         filepath = self._to_filepath('__metadata__')
         if os.path.isfile(filepath):
             meta = read_csv(filepath, wide=False)

larray/inout/excel.py

@@ -12,6 +12,8 @@
 except ImportError:
     xlsxwriter = None
 
+from typing import List, Tuple
+
 from larray.core.array import Array, asarray
 from larray.core.constants import nan
 from larray.core.group import _translate_sheet_name
@@ -243,10 +245,10 @@ def _open_for_write(self):
         engine = 'xlsxwriter' if ext == '.xlsx' and xlsxwriter is not None else None
         self.handle = pd.ExcelWriter(self.fname, engine=engine)
 
-    def list_items(self):
+    def list_items(self) -> List[Tuple[str, str]]:
         return [(name, 'Array') for name in self.handle.sheet_names if name != '__metadata__']
 
-    def _read_item(self, key, type, *args, **kwargs):
+    def _read_item(self, key, type, *args, **kwargs) -> Array:
         if type == 'Array':
             df = self.handle.parse(key, *args, **kwargs)
             return df_asarray(df, raw=True)
@@ -260,7 +262,7 @@ def _dump_item(self, key, value, *args, **kwargs):
         else:
             raise TypeError()
 
-    def _read_metadata(self):
+    def _read_metadata(self) -> Metadata:
         sheet_meta = '__metadata__'
         if sheet_meta in self.handle.sheet_names:
             meta = read_excel(self.handle, sheet_meta, engine=self.handle.engine, wide=False)
@@ -298,10 +300,10 @@ def _open_for_read(self):
     def _open_for_write(self):
         self.handle = open_excel(self.fname, overwrite_file=self.overwrite_file)
 
-    def list_items(self):
+    def list_items(self) -> List[Tuple[str, str]]:
         return [(name, 'Array') for name in self.handle.sheet_names() if name != '__metadata__']
 
-    def _read_item(self, key, type, *args, **kwargs):
+    def _read_item(self, key, type, *args, **kwargs) -> Array:
         if type == 'Array':
             return self.handle[key].load(*args, **kwargs)
         else:
@@ -313,7 +315,7 @@ def _dump_item(self, key, value, *args, **kwargs):
         else:
             raise TypeError()
 
-    def _read_metadata(self):
+    def _read_metadata(self) -> Metadata:
         sheet_meta = '__metadata__'
         if sheet_meta in self.handle:
             meta = self.handle[sheet_meta].load(wide=False)

larray/inout/hdf.py

@@ -4,12 +4,15 @@
 import pandas as pd
 from pandas import HDFStore
 
+from typing import List, Tuple, Union
+
 from larray.core.array import Array
 from larray.core.axis import Axis
 from larray.core.constants import nan
 from larray.core.group import Group, LGroup, _translate_group_key_hdf
 from larray.core.metadata import Metadata
 from larray.util.misc import LHDFStore
+from larray.util.types import Scalar
 from larray.inout.session import register_file_handler
 from larray.inout.common import FileHandler, _supported_typenames, _supported_scalars_types
 from larray.inout.pandas import df_asarray
@@ -136,7 +139,7 @@ def _open_for_read(self):
     def _open_for_write(self):
         self.handle = HDFStore(self.fname)
 
-    def list_items(self):
+    def list_items(self) -> List[Tuple[str, str]]:
         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) ----
@@ -146,7 +149,7 @@ def list_items(self):
         items += [(key.split('/')[-1], 'Group_Backward_Comp') for key in keys if '__groups__' in key]
         return items
 
-    def _read_item(self, key, typename, *args, **kwargs):
+    def _read_item(self, key, typename, *args, **kwargs) -> Union[Array, Axis, Group, Scalar]:
         if typename in _supported_typenames:
             hdf_key = '/' + key
         # ---- for backward compatibility (LArray < 0.33) ----
@@ -172,7 +175,7 @@ def _dump_item(self, key, value, *args, **kwargs):
         else:
             raise TypeError()
 
-    def _read_metadata(self):
+    def _read_metadata(self) -> Metadata:
         metadata = Metadata.from_hdf(self.handle)
         if metadata is None:
             metadata = Metadata()

larray/inout/pickle.py

@@ -2,12 +2,15 @@
 import os.path
 from collections import OrderedDict
 
+from typing import List, Tuple, Union
+
 from larray.core.axis import Axis
 from larray.core.group import Group
 from larray.core.array import Array
 from larray.core.metadata import Metadata
 from larray.inout.session import register_file_handler
 from larray.inout.common import FileHandler, _supported_types, _supported_typenames, _supported_scalars_types
+from larray.util.types import Scalar
 
 
 @register_file_handler('pickle', ['pkl', 'pickle'])
@@ -22,7 +25,7 @@ def _open_for_write(self):
         else:
             self.data = OrderedDict()
 
-    def list_items(self):
+    def list_items(self) -> List[Tuple[str, str]]:
         # scalar
         items = [(key, type(value).__name__) for key, value in self.data.items()
                  if isinstance(value, _supported_scalars_types)]
@@ -34,7 +37,7 @@ def list_items(self):
         items += [(key, 'Array') for key, value in self.data.items() if isinstance(value, Array)]
         return items
 
-    def _read_item(self, key, typename, *args, **kwargs):
+    def _read_item(self, key, typename, *args, **kwargs) -> Union[Array, Axis, Group, Scalar]:
         if typename in _supported_typenames:
             return self.data[key]
         else:
@@ -46,7 +49,7 @@ def _dump_item(self, key, value, *args, **kwargs):
         else:
             raise TypeError()
 
-    def _read_metadata(self):
+    def _read_metadata(self) -> Metadata:
         if '__metadata__' in self.data:
             return self.data['__metadata__']
         else:

larray/inout/session.py

@@ -1,3 +1,5 @@
+from larray.inout.common import FileHandler
+
 
 handler_classes = {}
 ext_default_engine = {}
@@ -13,7 +15,7 @@ def register_file_handler(engine, extensions=None):
     extensions : str or list of str, optional
         Extension(s) associated with the file handler.
     """
-    def decorate_class(cls):
+    def decorate_class(cls: FileHandler) -> FileHandler:
         if engine not in handler_classes:
             handler_classes[engine] = cls
         if extensions is None:
@@ -28,7 +30,7 @@ def decorate_class(cls):
     return decorate_class
 
 
-def get_file_handler(engine):
+def get_file_handler(engine) -> FileHandler:
     if engine not in handler_classes:
         raise TypeError(f"Engine {engine} is currently not implemented")
     file_handler_cls = handler_classes[engine]