Merge pull request pandas-dev#56 from bmoscon/master

new VersionStore API - get_data_info

Author: bmoscon

CHANGES.md

@@ -1,6 +1,10 @@
 
 ## Changelog
 
+### 1.15 (2015-11-25)
+
+  * Feature: get_data_info API added to version_store. 
+  
 ### 1.14 (2015-11-25)
 ### 1.12 (2015-11-12)
 

arctic/store/_ndarray_store.py

@@ -150,27 +150,17 @@ def _index_range(self, version, symbol, from_version=None, **kwargs):
             from_index = from_version['up_to']
         return from_index, None
 
-    def get_info(self, arctic_lib, version, symbol, **kwargs):
-        collection = arctic_lib.get_top_level_collection()
+    def get_info(self, version):
+        ret = {}
         dtype = self._dtype(version['dtype'], version.get('dtype_metadata', {}))
         length = int(version['up_to'])
-
-        spec = {'symbol': symbol,
-                'parent': version.get('base_version_id', version['_id']),
-                'segment': {'$lt': length}}
-
-        n_segments = collection.find(spec).count()
-
-        est_size = dtype.itemsize * length
-        return """Handler: %s
-
-dtype: %s
-
-%d rows in %d segments
-Data size: %s bytes
-
-Version document:
-%s""" % (self.__class__.__name__, dtype, length, n_segments, est_size, pprint.pformat(version))
+        ret['size'] = dtype.itemsize * length
+        ret['segment_count'] = version['segment_count']
+        ret['dtype'] = version['dtype']
+        ret['type'] = version['type']
+        ret['handler'] = self.__class__.__name__
+        ret['rows'] = int(version['up_to'])
+        return ret
 
     def read(self, arctic_lib, version, symbol, read_preference=None, **kwargs):
         index_range = self._index_range(version, symbol, **kwargs)

arctic/store/_pandas_ndarray_store.py

@@ -4,6 +4,7 @@
 from pandas import DataFrame, MultiIndex, Series, DatetimeIndex, Panel
 from pandas.tslib import Timestamp, get_timezone
 import numpy as np
+import ast
 
 from .._compression import compress, decompress
 from ..exceptions import ArcticException
@@ -195,6 +196,17 @@ def read(self, arctic_lib, version, symbol, read_preference=None, date_range=Non
             item = self._daterange(item, date_range)
         return item
 
+    def get_info(self, version):
+        """
+        parses out the relevant information in version
+        and returns it to the user in a dictionary
+        """
+        ret = super(PandasStore, self).get_info(version)
+        ret['col_names'] = version['dtype_metadata']
+        ret['handler'] = self.__class__.__name__
+        ret['dtype'] = ast.literal_eval(version['dtype'])
+        return ret
+
 
 def _start_end(date_range, dts):
     """

arctic/store/_pickle_store.py

@@ -20,12 +20,11 @@ class PickleStore(object):
     def initialize_library(cls, *args, **kwargs):
         pass
 
-    def get_info(self, arctic_lib, version, symbol, **kwargs):
-        if 'blob' in version:
-            if version['blob'] != _MAGIC_CHUNKED:
-                version['blob'] = "<Compressed pickle.....>"
-
-        return """Handler: %s\n\nVersion document:\n%s""" % (self.__class__.__name__, pprint.pformat(version))
+    def get_info(self, version):
+        ret = {}
+        ret['type'] = 'blob'
+        ret['handler'] = self.__class__.__name__
+        return ret
 
     def read(self, mongoose_lib, version, symbol, **kwargs):
         blob = version.get("blob")

arctic/store/version_store.py

@@ -335,10 +335,9 @@ def read(self, symbol, as_of=None, date_range=None, from_version=None, allow_sec
             raise
 
     @mongo_retry
-    def _show_info(self, symbol, as_of=None):
+    def get_info(self, symbol, as_of=None):
         """
-        Print details on the stored symbol: the underlying storage handler
-        and the version_document corresponding to the specified version.
+        Reads and returns information about the data stored for symbol
 
         Parameters
         ----------
@@ -349,16 +348,19 @@ def _show_info(self, symbol, as_of=None):
             `int` : specific version number
             `str` : snapshot name which contains the version
             `datetime.datetime` : the version of the data that existed as_of the requested point in time
+
+        Returns
+        -------
+        dictionary of the information (specific to the type of data)
         """
-        print self._get_info(symbol, as_of)
+        version = self._read_metadata(symbol, as_of=as_of, read_preference=None)
+        print version
+        handler = self._read_handler(version, symbol)
+        if handler and hasattr(handler, 'get_info'):
+            return handler.get_info(version)
+        return {}
+
 
-    def _get_info(self, symbol, as_of=None):
-        _version = self._read_metadata(symbol, as_of=as_of)
-        handler = self._read_handler(_version, symbol)
-        if hasattr(handler, "get_info"):
-            return handler.get_info(self._arctic_lib, _version, symbol)
-        else:
-            return """Handler: %s\n\nVersion document:\n%s""" % (handler.__class__.__name__, pprint.pformat(_version))
 
     def _do_read(self, symbol, version, from_version=None, **kwargs):
         handler = self._read_handler(version, symbol)

tests/integration/store/test_ndarray_store.py

@@ -79,7 +79,7 @@ def test_save_read_big_2darray(library):
 def test_get_info_bson_object(library):
     ndarr = np.ones(1000)
     library.write('MYARR', ndarr)
-    assert library._get_info('MYARR').startswith('''Handler: NdarrayStore''')
+    assert library.get_info('MYARR')['handler'] == 'NdarrayStore'
 
 
 def test_save_read_ndarray_with_array_field(library):

tests/integration/store/test_pandas_store.py

@@ -16,6 +16,7 @@
 from arctic.store.version_store import register_versioned_storage
 from pandas.tseries.offsets import DateOffset
 
+
 register_versioned_storage(PandasDataFrameStore)
 
 
@@ -797,3 +798,42 @@ def test_daterange_fails_with_timezone_start(library):
     with pytest.raises(ValueError):
         library.read('MYARR', date_range=DateRange(start=dt(2015, 1, 1, tzinfo=mktz())))
 
+
+def test_data_info_series(library):
+    s = Series(data=[1, 2, 3], index=[4, 5, 6])
+    library.write('pandas', s)
+    md = library.get_info('pandas')
+    assert md == {'dtype': [('index', '<i8'), ('values', '<i8')],
+                  'col_names': {u'index': [u'index'], u'columns': [u'values']},
+                  'type': u'pandasseries',
+                  'handler': 'PandasSeriesStore',
+                  'rows': 3,
+                  'segment_count': 1,
+                  'size': 48}
+
+
+def test_data_info_df(library):
+    s = DataFrame(data=[1, 2, 3], index=[4, 5, 6])
+    library.write('pandas', s)
+    md = library.get_info('pandas')
+    assert md == {'dtype': [('index', '<i8'), ('0', '<i8')],
+                  'col_names': {u'index': [u'index'], u'columns': [u'0']},
+                  'type': u'pandasdf',
+                  'handler': 'PandasDataFrameStore',
+                  'rows': 3,
+                  'segment_count': 1,
+                  'size': 48}
+
+
+def test_data_info_cols(library):
+    i = MultiIndex.from_tuples([(1, "ab"), (2, "bb"), (3, "cb")])
+    s = DataFrame(data=[100, 200, 300], index=i)
+    library.write('test_data', s)
+    md = library.get_info('test_data')
+    assert md == {'dtype': [('level_0', '<i8'), ('level_1', 'S2'), ('0', '<i8')],
+                  'col_names': {u'index': [u'level_0', u'level_1'], u'columns': [u'0']},
+                  'type': u'pandasdf',
+                  'handler': 'PandasDataFrameStore',
+                  'rows': 3,
+                  'segment_count': 1,
+                  'size': 54}

tests/integration/store/test_pickle_store.py

@@ -31,7 +31,7 @@ def test_save_read_bson_object(library):
 def test_get_info_bson_object(library):
     blob = {'foo': dt(2015, 1, 1), 'object': Arctic}
     library.write('BLOB', blob)
-    assert library._get_info('BLOB').startswith('Handler: PickleStore')
+    assert library.get_info('BLOB')['handler'] == 'PickleStore'
 
 
 def test_bson_large_object(library):