def push_get1(comp):
buf = io.BytesIO(b'')
with proio.Writer(fileobj=buf) as writer:
writer.set_compression(comp)
eventsOut = []
event = proio.Event()
event.add_entries('MCParticle', prolcio.MCParticle(),
prolcio.MCParticle())
event.add_entries('TrackerHits', prolcio.SimTrackerHit(),
prolcio.SimTrackerHit())
writer.push(event)
eventsOut.append(event)
event = proio.Event()
event.add_entries('TrackerHits', prolcio.SimTrackerHit(),
prolcio.SimTrackerHit())
writer.push(event)
eventsOut.append(event)
buf.seek(0, 0)
with proio.Reader(fileobj=buf) as reader:
for i in range(0, len(eventsOut)):
event = reader.__next__()
assert event != None
assert event.__str__() == eventsOut[i].__str__()
def push_seek_skip_get1(comp):
buf = io.BytesIO(b'')
with proio.Writer(fileobj = buf) as writer:
writer.set_compression(comp)
eventsOut = []
event = proio.Event()
event.add_entries(
'Hit1',
prolcio.MCParticle(),
prolcio.MCParticle()
)
writer.push(event)
eventsOut.append(event)
event = proio.Event()
event.add_entries(
'Hit2',
prolcio.SimTrackerHit(),
prolcio.SimTrackerHit()
)
writer.push(event)
eventsOut.append(event)
writer.flush()
event = proio.Event()
event.add_entries(
'Hit3',
prolcio.SimTrackerHit(),
prolcio.SimTrackerHit()
)
writer.push(event)
eventsOut.append(event)
buf.seek(0, 0)
with proio.Reader(fileobj = buf) as reader:
event = reader.next()
assert event != None
assert event.__str__() == eventsOut[0].__str__()
reader.seek_to_start()
reader.skip(2)
event = reader.next()
assert event != None
assert event.__str__() == eventsOut[2].__str__()
def test_ref_deref2():
event = proio.Event()
parent = prolcio.MCParticle()
parent.PDG = 443
parent_id = event.add_entry('MCParticles', parent)
child1 = prolcio.MCParticle()
child1.PDG = 11
child2 = prolcio.MCParticle()
child2.PDG = -11
child_ids = event.add_entries('MCParticles', child1, child2)
parent.children.extend(child_ids)
child1.parents.append(parent_id)
child2.parents.append(parent_id)
mc_particles = event.tagged_entries('MCParticles')
assert mc_particles != None
parent_ = event.get_entry(mc_particles[0])
assert parent_.__str__() == parent.__str__()
child1_ = event.get_entry(parent_.children[0])
assert child1_.__str__() == child1.__str__()
child2_ = event.get_entry(parent_.children[1])
assert child2_.__str__() == child2.__str__()
parent_ = event.get_entry(child1_.parents[0])
assert parent_.__str__() == parent.__str__()
parent_ = event.get_entry(child2_.parents[0])
assert parent_.__str__() == parent.__str__()
def test_tag_cleanup1():
event = proio.Event()
entry = eic.SimHit()
entryID = event.add_entry('SimCreated', entry)
assert (len(event.tagged_entries('SimCreated')) == 1)
assert (len(event.all_entries()) == 1)
event.remove_entry(entryID)
assert (len(event.tagged_entries('SimCreated')) == 0)
assert (len(event.all_entries()) == 0)
def write_fifo(filename):
with proio.Writer(filename) as writer:
event = proio.Event()
writer.push(event)
writer.flush()
writer.push(event)
writer.push(event)
writer.flush()
writer.push(event)
writer.push(event)
def test_tag_delete1():
event = proio.Event()
entry = eic.SimHit()
entryID = event.add_entry('SimCreated', entry)
assert (len(event.tags()) == 1)
assert (len(event.all_entries()) == 1)
event.delete_tag('WrongTag')
assert (len(event.tags()) == 1)
event.delete_tag('SimCreated')
assert (len(event.tags()) == 0)
assert (len(event.all_entries()) == 1)
def test_untag2():
event = proio.Event()
entry = eic.SimHit()
entryID = event.add_entry('SimCreated', entry)
assert (len(event.tagged_entries('SimCreated')) == 1)
event.untag_entry(entryID + 1, 'WrongTag')
assert (len(event.tagged_entries('SimCreated')) == 1)
assert (len(event.tags()) == 1)
assert (len(event.all_entries()) == 1)
event.untag_entry(entryID + 1, 'SimCreated')
assert (len(event.tagged_entries('SimCreated')) == 1)
assert (len(event.all_entries()) == 1)
def test_tag_rev_lookup1():
event = proio.Event()
entry = eic.SimHit()
entryID = event.add_entry('SimCreated', entry)
assert (len(event.entry_tags(entryID)) == 1)
assert (event.entry_tags(entryID)[0] == 'SimCreated')
assert (len(event.tags()) == 1)
event.tag_entry(entryID, 'Tracker')
assert (len(event.entry_tags(entryID)) == 2)
assert (len(event.tags()) == 2)
event.untag_entry(entryID, 'SimCreated')
assert (len(event.entry_tags(entryID)) == 1)
assert (event.entry_tags(entryID)[0] == 'Tracker')
assert (len(event.tags()) == 2)
def test_header1():
buf = io.BytesIO(b'')
with proio.Writer(fileobj = buf) as writer:
event = proio.Event()
writer.push(event)
writer.push(event)
writer.flush()
writer.push(event)
buf.seek(0, 0)
with proio.Reader(fileobj = buf) as reader:
reader.skip(0)
assert(reader._bucket_header.nEvents == 2)
reader.skip(2)
assert(reader._bucket_header.nEvents == 1)
reader.skip(1)
assert(reader._bucket_header is None)
def push_get4(comp):
buf = io.BytesIO(b'')
with proio.Writer(fileobj=buf) as writer:
writer.set_compression(comp)
eventsOut = []
event = proio.Event()
event.add_entry('test', descriptor_pb2.FileDescriptorProto())
writer.push(event)
eventsOut.append(event)
buf.seek(0, 0)
with proio.Reader(fileobj=buf) as reader:
for i in range(0, len(eventsOut)):
event = reader.__next__()
assert event != None
assert event.__str__() == eventsOut[i].__str__()
def test_write_iterate_file():
n_events = 5
test_file = tempfile.NamedTemporaryFile()
with proio.Writer(fileobj=test_file) as writer:
event = proio.Event()
for i in range(0, n_events):
writer.push(event)
n_events = 0
test_file.seek(0)
with proio.Reader(fileobj=test_file) as reader:
for event in reader:
n_events += 1
test_file.close()
assert (n_events == 5)
def test_push_update1():
buf = io.BytesIO(b'')
with proio.Writer(fileobj = buf) as writer:
writer.push_metadata('key1', b'value1')
writer.push_metadata('key2', b'value2')
event = proio.Event()
writer.push(event)
writer.push_metadata('key2', b'value3')
writer.push(event)
writer.push_metadata('key1', b'value4')
writer.push_metadata('key2', b'value5')
writer.push(event)
buf.seek(0, 0)
with proio.Reader(fileobj = buf) as reader:
event1 = next(reader)
event2 = next(reader)
event3 = next(reader)
assert(event1.metadata['key1'] == b'value1')
assert(event1.metadata['key2'] == b'value2')
assert(event2.metadata['key1'] == b'value1')
assert(event2.metadata['key2'] == b'value3')
assert(event3.metadata['key1'] == b'value4')
assert(event3.metadata['key2'] == b'value5')
buf = io.BytesIO(b'')
with proio.Writer(fileobj = buf) as writer:
writer.push(event1)
writer.push(event2)
writer.push(event3)
buf.seek(0, 0)
with proio.Reader(fileobj = buf) as reader:
event1 = next(reader)
event2 = next(reader)
event3 = next(reader)
assert(event1.metadata['key1'] == b'value1')
assert(event1.metadata['key2'] == b'value2')
assert(event2.metadata['key1'] == b'value1')
assert(event2.metadata['key2'] == b'value3')
assert(event3.metadata['key1'] == b'value4')
assert(event3.metadata['key2'] == b'value5')
def test_next_header1():
buf = io.BytesIO(b'')
with proio.Writer(fileobj = buf) as writer:
event = proio.Event()
writer.push(event)
writer.push(event)
writer.flush()
writer.push(event)
buf.seek(0, 0)
with proio.Reader(fileobj = buf) as reader:
header = reader.next_header()
assert(header.nEvents == 2)
assert(reader.next() != None)
try:
reader.next()
assert(False)
except StopIteration:
pass
def test_write_seek_file():
test_file = tempfile.NamedTemporaryFile()
with proio.Writer(fileobj=test_file) as writer:
event = proio.Event()
writer.push(event)
writer.flush()
writer.push(event)
writer.push(event)
writer.flush()
writer.push(event)
writer.push(event)
test_file.seek(0)
with proio.Reader(fileobj=test_file) as reader:
n_events = 4
reader.skip(n_events)
for event in reader:
n_events += 1
assert (n_events == 5)
test_file.close()
def test_ref_deref1():
buf = io.BytesIO(b'')
with proio.Writer(fileobj=buf) as writer:
event = proio.Event()
parent = prolcio.MCParticle()
parent.PDG = 443
parent_id = event.add_entry('MCParticles', parent)
child1 = prolcio.MCParticle()
child1.PDG = 11
child2 = prolcio.MCParticle()
child2.PDG = -11
child_ids = event.add_entries('MCParticles', child1, child2)
parent.children.extend(child_ids)
child1.parents.append(parent_id)
child2.parents.append(parent_id)
writer.push(event)
buf.seek(0, 0)
with proio.Reader(fileobj=buf) as reader:
event = reader.next()
assert event != None
mc_particles = event.tagged_entries('MCParticles')
assert mc_particles != None
parent_ = event.get_entry(mc_particles[0])
assert parent_.__str__() == parent.__str__()
child1_ = event.get_entry(parent_.children[0])
assert child1_.__str__() == child1.__str__()
child2_ = event.get_entry(parent_.children[1])
assert child2_.__str__() == child2.__str__()
parent_ = event.get_entry(child1_.parents[0])
assert parent_.__str__() == parent.__str__()
parent_ = event.get_entry(child2_.parents[0])
assert parent_.__str__() == parent.__str__()
def test_get_bad_entry1():
event = proio.Event()
assert (event.get_entry(1) is None)
