Skip to content

Commit

Permalink
Merge pull request pandas-dev#4 in DATA/arctic from pass_parameters_t…
Browse files Browse the repository at this point in the history
…hrough_to_tickstore to master

* commit 'f3d7229c98cde50257394dddb41400c44ea051b0':
  MDP-448 Pass all arguments through from TopLevelTickStore.read to TickStore.read
  • Loading branch information
jamesblackburn committed Aug 19, 2015
2 parents 741bda5 + f3d7229 commit c7a5c8d
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 3 deletions.
5 changes: 3 additions & 2 deletions arctic/tickstore/toplevel.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,9 +98,10 @@ def add(self, date_range, library_name):
self._collection.update_one({'library_name': library_name},
{'$set': {'start': start, 'end': end}}, upsert=True)

def read(self, symbol, date_range, columns=['BID', 'ASK', 'TRDPRC_1', 'BIDSIZE', 'ASKSIZE', 'TRDVOL_1'], **kwargs):
def read(self, symbol, date_range, columns=['BID', 'ASK', 'TRDPRC_1', 'BIDSIZE', 'ASKSIZE', 'TRDVOL_1'], include_images=False):
libraries = self._get_libraries(date_range)
dfs = [l.library.read(symbol, l.date_range.intersection(date_range), columns) for l in libraries]
dfs = [l.library.read(symbol, l.date_range.intersection(date_range), columns,
include_images=include_images) for l in libraries]
return pd.concat(dfs)

def write(self, symbol, data):
Expand Down
21 changes: 20 additions & 1 deletion tests/unit/tickstore/test_toplevel.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from mock import Mock, patch, MagicMock, create_autospec, sentinel
from mock import Mock, patch, MagicMock, create_autospec, sentinel, call
import pytest
from datetime import datetime as dt
import pandas as pd
Expand All @@ -10,6 +10,7 @@
from arctic.exceptions import OverlappingDataException
from arctic.tickstore.toplevel import TopLevelTickStore, TickStoreLibrary
from dateutil.rrule import rrule, DAILY
from arctic.tickstore.tickstore import TickStore


def test_raise_exception_if_daterange_is_not_provided():
Expand Down Expand Up @@ -145,3 +146,21 @@ def test_write_pandas_data_to_right_libraries():
TopLevelTickStore.write(self, 'blah', sentinel.data)
mock_lib1.write.assert_called_once_with('blah', slice1)
mock_lib2.write.assert_called_once_with('blah', slice2)


def test_read():
self = create_autospec(TopLevelTickStore)
tsl = TickStoreLibrary(create_autospec(TickStore), create_autospec(DateRange))
self._get_libraries.return_value = [tsl, tsl]
dr = create_autospec(DateRange)
with patch('pandas.concat') as concat:
res = TopLevelTickStore.read(self, sentinel.symbol, dr,
columns=sentinel.include_columns,
include_images=sentinel.include_images)
assert concat.call_args_list == [call([tsl.library.read.return_value,
tsl.library.read.return_value])]
assert res == concat.return_value
assert tsl.library.read.call_args_list == [call(sentinel.symbol, tsl.date_range.intersection.return_value,
sentinel.include_columns, include_images=sentinel.include_images),
call(sentinel.symbol, tsl.date_range.intersection.return_value,
sentinel.include_columns, include_images=sentinel.include_images)]

0 comments on commit c7a5c8d

Please sign in to comment.