def test_pivot_table(self):
df = DataFrame(ohlc)
df = df.pivot_table(index=['date', 'tick'], values=['price', 'qty'], columns=['time'], fill_value=0)
self.assertIn('price_open', df)
self.assertIn('price_high', df)
self.assertIn('qty_open', df)
self.assertIn('qty_high', df)
self.assertEqual(len(df), 18)
# Fill missing
missing_price = df[(df.date == '2019-01-02') & (df.tick == 'goog')].price_close[0]
self.assertEqual(missing_price, 0)
def test_pivot(self):
df = DataFrame(tickers)
df['price'] = df.price.apply(float)
self.assertEqual(len(df), 18)
self.assertNotIn('aapl', df)
df = df.pivot(index='date', columns='tick', values='price')
self.assertEqual(len(df), 6)
self.assertIn('aapl', df)
self.assertIn('goog', df)
self.assertIn('msft', df)
self.assertEqual(100, df['aapl'][0])
self.assertEqual(100.10, df['aapl'][1])
def test_creation(self):
data = [list(x.values()) for x in self.df.iterrows()]
df = DataFrame(values=data, columns=list(tickers.keys()))
self.assertEqual(
self.df.tick[5],
df.tick[5],
)
self.assertEqual(
self.df.price[7],
df.price[7],
)
def setUp(self):
self.df = DataFrame(tickers)
def setUp(self):
self.df = DataFrame(tickers)
self.df['date'] = self.df['date'].apply(str_to_dt)
def setUp(self):
self.df = DataFrame(tickers)
self.pdf = DataFrame(positions)
class TestDataFrameSugar(unittest.TestCase):
def setUp(self):
self.df = DataFrame(tickers)
self.pdf = DataFrame(positions)
def test_get(self):
# Dict access
items = self.df['price']
self.assertIsInstance(items, Series)
for price, expected in zip(items, expected_result):
self.assertEqual(str(price), str(expected['price']))
# Select multiple columns
items = self.df[['price', 'tick']]
for price, expected in zip(items.price, expected_result):
self.assertEqual(str(price), str(expected['price']))
for tick, expected in zip(items.tick, expected_result):
self.assertEqual(str(tick), str(expected['tick']))
# Attribute access
items = self.df.tick
self.assertIsInstance(items, Series)
for tick, expected in zip(items, expected_result):
self.assertEqual(tick, expected['tick'])
def test_set(self):
# Dict setting
self.df['position'] = None
for item in self.df['position']:
self.assertIsNone(item)
# Set a list
self.df['position'] = ['open']*len(self.df)
for item in self.df['position']:
self.assertEqual(item, 'open')
def test_filter(self):
results = self.df['date'] == '2019-01-02'
self.assertListEqual(
[False]*3 + [True]*3 + [False]*12,
list(results)
)
results = self.df['date'] >= '2019-01-02'
self.assertListEqual(
[False]*3 + [True]*15,
list(results)
)
self.assertEqual(18, len(self.df))
df = self.df[self.df['date'] >= '2019-01-03']
self.assertEqual(12, len(df))
self.assertListEqual(
['2019-01-03', '2019-01-03', '2019-01-03','2019-01-04', '2019-01-04', '2019-01-04','2019-01-05', '2019-01-05', '2019-01-05','2019-01-06', '2019-01-06', '2019-01-06'],
list(df['date']),
)
df = self.df[(self.df['date'] >= '2019-01-02') & (self.df['date'] <= '2019-01-04')]
self.assertEqual(9, len(df))
self.assertListEqual(
['2019-01-02', '2019-01-02', '2019-01-02','2019-01-03', '2019-01-03', '2019-01-03','2019-01-04', '2019-01-04', '2019-01-04'],
list(df['date']),
)
def test_drop(self):
results = self.df['date'] == '2019-01-02'
self.assertEqual(18, len(self.df))
self.df.drop(results)
self.assertEqual(15, len(self.df))
results = self.df['date'] == '2019-01-02'
self.assertEqual(0, len(self.df[results]))
results = self.df['date'] == '2019-01-01'
self.assertEqual(3, len(self.df[results]))
results = self.df['date'] == '2019-01-03'
self.assertEqual(3, len(self.df[results]))
results = self.df['date'] == '2019-01-04'
self.assertEqual(3, len(self.df[results]))
results = self.df['date'] == '2019-01-05'
self.assertEqual(3, len(self.df[results]))
results = self.df['date'] == '2019-01-06'
self.assertEqual(3, len(self.df[results]))
def test_drop_all(self):
self.df.drop('all')
self.assertEqual(0, len(self.df))
def test_contains(self):
self.assertIn('date', self.df)
self.assertNotIn('season', self.df)
def test_features(self):
self.df['date'] = self.df['date'].apply(str_to_dt)
self.df['price'] = self.df['price'].apply(float)
self.pdf['date'] = self.pdf['date'].apply(str_to_dt)
self.pdf['position'] = self.pdf['position'].apply(lambda x: x.strip())
for row in self.pdf.iterrows():
sset = (self.df['date'] >= row['date']) & (self.df['tick'] == row['tick'])
self.df.set(sset, 'position', row['position'])
for actual, expected in zip(self.df.iterrows(), expected_result):
self.assertEqual(actual['tick'], expected['tick'])
self.assertEqual(actual['price'], expected['price'])
self.assertEqual(actual['position'], expected['position'])
class TestDataFrame(unittest.TestCase):
def setUp(self):
self.df = DataFrame(tickers)
self.pdf = DataFrame(positions)
def test_creation(self):
data = [list(x.values()) for x in self.df.iterrows()]
df = DataFrame(values=data, columns=list(tickers.keys()))
self.assertEqual(
self.df.tick[5],
df.tick[5],
)
self.assertEqual(
self.df.price[7],
df.price[7],
)
def test_apply(self):
for date in self.df.get('date'):
self.assertIsInstance(date, str)
self.df['date'] = self.df['date'].apply(str_to_dt)
for date in self.df['date']:
self.assertIsInstance(date, dt.datetime)
def test_get(self):
items = self.df.get('price')
for price, expected in zip(items, expected_result):
self.assertEqual(str(price), str(expected['price']))
items = self.df.get('space')
self.assertListEqual(
[None]*18,
list(items)
)
# Doesnt exist
items = self.df.get('space', 0)
self.assertListEqual(
[0]*18,
list(items)
)
# Select multiple columns
items = self.df.get(['price', 'tick'])
for price, expected in zip(items.price, expected_result):
self.assertEqual(str(price), str(expected['price']))
for tick, expected in zip(items.tick, expected_result):
self.assertEqual(str(tick), str(expected['tick']))
def test_set(self):
self.df.set([True]*len(self.df), 'position', None)
for item in self.df.get('position'):
self.assertIsNone(item)
self.df.set('all', 'position1', None)
for item in self.df.get('position1'):
self.assertIsNone(item)
try:
self.df.unknown = 2
self.fail('Should not be able to set attributes')
except AttributeError:
pass
def test_iterrows(self):
rows = list(self.pdf.iterrows())
self.assertEqual(12, len(rows))
self.assertEqual('aapl', rows[0]['tick'])
self.assertEqual('msft', rows[5]['tick'])
def test_sum(self):
self.df['price'] = self.df.price.apply(float)
self.assertEqual(1631.12, sum(self.df.get('price')))