mirror of
https://github.com/wassname/catalyst.git
synced 2026-08-11 11:16:15 +08:00
Merge pull request #924 from quantopian/dataset-subclassing
ENH: Make datasets have subclass relationships
This commit is contained in:
@@ -90,7 +90,7 @@ class BlazeToPipelineTestCase(TestCase):
|
||||
self.assertEqual(ds.__name__, name)
|
||||
self.assertTrue(issubclass(ds, DataSet))
|
||||
self.assertEqual(
|
||||
{c.name: c.dtype for c in ds._columns},
|
||||
{c.name: c.dtype for c in ds.columns},
|
||||
{'sid': np.int64, 'value': np.float64},
|
||||
)
|
||||
|
||||
|
||||
@@ -28,6 +28,14 @@ class SomeDataSet(DataSet):
|
||||
buzz = Column(float64_dtype)
|
||||
|
||||
|
||||
class SubDataSet(SomeDataSet):
|
||||
pass
|
||||
|
||||
|
||||
class SubDataSetNewCol(SomeDataSet):
|
||||
qux = Column(float64_dtype)
|
||||
|
||||
|
||||
class SomeFactor(Factor):
|
||||
dtype = float64_dtype
|
||||
window_length = 5
|
||||
@@ -321,3 +329,66 @@ class ObjectIdentityTestCase(TestCase):
|
||||
|
||||
with self.assertRaises(InvalidDType):
|
||||
SomeFactor(dtype=1)
|
||||
|
||||
|
||||
class SubDataSetTestCase(TestCase):
|
||||
def test_subdataset(self):
|
||||
some_dataset_map = {
|
||||
column.name: column for column in SomeDataSet.columns
|
||||
}
|
||||
sub_dataset_map = {
|
||||
column.name: column for column in SubDataSet.columns
|
||||
}
|
||||
self.assertEqual(
|
||||
{column.name for column in SomeDataSet.columns},
|
||||
{column.name for column in SubDataSet.columns},
|
||||
)
|
||||
for k, some_dataset_column in some_dataset_map.items():
|
||||
sub_dataset_column = sub_dataset_map[k]
|
||||
self.assertIsNot(
|
||||
some_dataset_column,
|
||||
sub_dataset_column,
|
||||
'subclass column %r should not have the same identity as'
|
||||
' the parent' % k,
|
||||
)
|
||||
self.assertEqual(
|
||||
some_dataset_column.dtype,
|
||||
sub_dataset_column.dtype,
|
||||
'subclass column %r should have the same dtype as the parent' %
|
||||
k,
|
||||
)
|
||||
|
||||
def test_add_column(self):
|
||||
some_dataset_map = {
|
||||
column.name: column for column in SomeDataSet.columns
|
||||
}
|
||||
sub_dataset_new_col_map = {
|
||||
column.name: column for column in SubDataSetNewCol.columns
|
||||
}
|
||||
sub_col_names = {column.name for column in SubDataSetNewCol.columns}
|
||||
|
||||
# check our extra col
|
||||
self.assertIn('qux', sub_col_names)
|
||||
self.assertEqual(
|
||||
sub_dataset_new_col_map['qux'].dtype,
|
||||
float64_dtype,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
{column.name for column in SomeDataSet.columns},
|
||||
sub_col_names - {'qux'},
|
||||
)
|
||||
for k, some_dataset_column in some_dataset_map.items():
|
||||
sub_dataset_column = sub_dataset_new_col_map[k]
|
||||
self.assertIsNot(
|
||||
some_dataset_column,
|
||||
sub_dataset_column,
|
||||
'subclass column %r should not have the same identity as'
|
||||
' the parent' % k,
|
||||
)
|
||||
self.assertEqual(
|
||||
some_dataset_column.dtype,
|
||||
sub_dataset_column.dtype,
|
||||
'subclass column %r should have the same dtype as the parent' %
|
||||
k,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user