def Reader(path, proto=None, compression_type=None):
"""A TFRecordReader that defaults to tf.Example protos."""
if not proto:
proto = example_pb2.Example
return genomics_reader.TFRecordReader(
path, proto, compression_type=compression_type)
def testCompressed(self):
reader = genomics_reader.TFRecordReader(
test_utils.genomics_core_testdata('test_features.gff.tfrecord.gz'),
gff_pb2.GffRecord(),
)
records = list(reader.iterate())
self.assertEqual('GenBank', records[0].source)
self.assertEqual('ctg123', records[1].range.reference_name)
def testTwoIteratorsAtTheSameTime(self):
dreader = genomics_reader.TFRecordReader('0,1,2,3,4,5', DummyProto())
iter2 = iter(dreader)
for i in range(6):
self.assertEqual(str(i), six.next(dreader))
self.assertEqual(str(i), six.next(iter2))
def testMock(self):
reader = genomics_reader.TFRecordReader('a,b,c,d,e', DummyProto())
self.assertEqual(['a', 'b', 'c', 'd', 'e'], list(reader))
