diff --git a/tests/test_assets.py b/tests/test_assets.py index 16205ede..fea726ab 100644 --- a/tests/test_assets.py +++ b/tests/test_assets.py @@ -20,7 +20,7 @@ Tests for the zipline.assets package import sys from unittest import TestCase -from zipline.assets._assets import Asset +from zipline.assets._assets import Asset, Future class AssetTestCase(TestCase): @@ -36,7 +36,7 @@ class AssetTestCase(TestCase): self.assertEquals(str(Asset(5061)), 'Asset(5061)') -class TestSecurityRichCmp(TestCase): +class TestAssetRichCmp(TestCase): def test_lt(self): self.assertTrue(Asset(3) < Asset(4)) @@ -72,3 +72,18 @@ class TestSecurityRichCmp(TestCase): Asset(3) < 'a' with self.assertRaises(TypeError): 'a' < Asset(3) + + +class testFuture(TestCase): + + def test_repr(self): + + future = Future(2468, + notice_date='2014-01-20', + expiration_date='2014-02-20') + rep = future.__repr__() + + self.assertTrue("Future" in rep) + self.assertTrue("2468" in rep) + self.assertTrue("notice_date='2014-01-20'" in rep) + self.assertTrue("expiration_date='2014-02-20'" in rep) diff --git a/zipline/assets/_assets.pyx b/zipline/assets/_assets.pyx index 5a8bb4ae..16b8ba06 100644 --- a/zipline/assets/_assets.pyx +++ b/zipline/assets/_assets.pyx @@ -21,6 +21,9 @@ cimport cython import numpy as np cimport numpy as np +cdef enum AssetType: + EQUITY = 1 + FUTURE = 2 cdef class Asset: @@ -30,6 +33,7 @@ cdef class Asset: cdef readonly object symbol cdef readonly object asset_name + cdef readonly AssetType asset_type # TODO: Maybe declare as pandas Timestamp? cdef readonly object start_date @@ -45,10 +49,12 @@ cdef class Asset: object start_date=None, object end_date=None, object first_traded=None, - object exchange=""): + object exchange="", + *args, + **kwargs): self.sid = sid - self.sid_hash = hash(sid) + self.sid_hash = hash(sid) self.symbol = symbol self.asset_name = asset_name self.exchange = exchange @@ -142,7 +148,7 @@ cdef class Asset: for attr in attrs) strings = ('%s=%s' % (t[0], t[1]) for t in tuples) params = ', '.join(strings) - return 'Security(%d, %s)' % (self.sid, params) + return 'Asset(%d, %s)' % (self.sid, params) cpdef __reduce__(self): """ @@ -179,3 +185,57 @@ cdef class Asset: Build an Asset instance from a dict. """ return Asset(**dict_) + + +cdef class Equity(Asset): + + def __cinit__(self, + int sid, # sid is required + object symbol="", + object asset_name="", + object start_date=None, + object end_date=None, + object first_traded=None, + object exchange=""): + + self.asset_type = EQUITY + + def __repr__(self): + attrs = ('symbol', 'asset_name', 'exchange', + 'start_date', 'end_date', 'first_traded') + tuples = ((attr, repr(getattr(self, attr, None))) + for attr in attrs) + strings = ('%s=%s' % (t[0], t[1]) for t in tuples) + params = ', '.join(strings) + return 'Equity(%d, %s)' % (self.sid, params) + + +cdef class Future(Asset): + + cdef readonly object notice_date + cdef readonly object expiration_date + + def __cinit__(self, + int sid, # sid is required + object symbol="", + object asset_name="", + object start_date=None, + object end_date=None, + object notice_date=None, + object expiration_date=None, + object first_traded=None, + object exchange=""): + + self.asset_type = FUTURE + self.notice_date = notice_date + self.expiration_date = expiration_date + + def __repr__(self): + attrs = ('symbol', 'asset_name', 'exchange', + 'start_date', 'end_date', 'first_traded', 'notice_date', + 'expiration_date') + tuples = ((attr, repr(getattr(self, attr, None))) + for attr in attrs) + strings = ('%s=%s' % (t[0], t[1]) for t in tuples) + params = ', '.join(strings) + return 'Future(%d, %s)' % (self.sid, params)