Skip to content

Commit

Permalink
Changed the event counters to increment by seq_num, not data. Changed…
Browse files Browse the repository at this point in the history
… seq_num to be a kwarg in ComposeEventPage
  • Loading branch information
evalott100 committed Jun 30, 2023
1 parent cd601d7 commit 2910dfe
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 6 deletions.
46 changes: 42 additions & 4 deletions event_model/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2197,6 +2197,20 @@ def compose_stop(
)(exit_status=exit_status, reason=reason, uid=uid, time=time, validate=validate)


def dict_of_lists_has_equal_size_lists(dictionary: Dict[str, List]) -> bool:
"""Return True if all lists are the same size in a Dict[str, List]."""

dictionary_values = iter(dictionary.values())
first_element_len = len(next(dictionary_values))
next_element = next(dictionary_values, None)

while next_element:
if len(next_element) != first_element_len:
return False
next_element = next(dictionary_values, None)
return True


@dataclass
class ComposeEventPage:
descriptor: EventDescriptor
Expand All @@ -2206,12 +2220,30 @@ def __call__(
self,
data: Dict[str, List],
timestamps: Dict[str, Any],
seq_num: List[int],
seq_num: Optional[List[int]] = None,
filled: Optional[Dict[str, List[Union[bool, str]]]] = None,
uid: Optional[List] = None,
time: Optional[List] = None,
validate: bool = True,
) -> EventPage:
assert dict_of_lists_has_equal_size_lists(timestamps), (
"Cannot compose event_page: event_page contains `timestamps` "
"list values of different lengths"
)
assert dict_of_lists_has_equal_size_lists(data), (
"Cannot compose event_page: event_page contains `data` "
"lists of different lengths"
)
assert len(next(iter(timestamps.values()))) == len(next(iter(data.values()))), (
"Cannot compose event_page: the lists in `timestamps` are of a different "
"length to those in `data`"
)

if seq_num is None:
last_seq_num = self.event_counters[self.descriptor["name"]]
seq_num = list(
range(last_seq_num, len(next(iter(data.values()))) + last_seq_num)
)
N = len(seq_num)
if uid is None:
uid = [str(uuid.uuid4()) for _ in range(N)]
Expand Down Expand Up @@ -2248,7 +2280,7 @@ def __call__(
"Keys in event['filled'] {} must be a subset of those in "
"event['data'] {}".format(filled.keys(), data.keys())
)
self.event_counters[self.descriptor["name"]] += len(data)
self.event_counters[self.descriptor["name"]] += len(seq_num)
return doc


Expand All @@ -2268,7 +2300,13 @@ def compose_event_page(
Here for backwards compatibility, the Compose class is prefered.
"""
return ComposeEventPage(descriptor, event_counters)(
data, timestamps, seq_num, filled, uid=uid, time=time, validate=validate
data,
timestamps,
seq_num=seq_num,
filled=filled,
uid=uid,
time=time,
validate=validate,
)


Expand Down Expand Up @@ -2324,7 +2362,7 @@ def __call__(
"Keys in event['filled'] {} must be a subset of those in "
"event['data'] {}".format(filled.keys(), data.keys())
)
self.event_counters[self.descriptor["name"]] += 1
self.event_counters[self.descriptor["name"]] = seq_num + 1
return doc


Expand Down
4 changes: 2 additions & 2 deletions event_model/tests/test_em.py
Original file line number Diff line number Diff line change
Expand Up @@ -1026,8 +1026,8 @@ def test_array_like():
)
desc_bundle.compose_event_page(
data={"a": dask_array.ones((5, 3))},
timestamps={"a": [1, 2, 3]},
seq_num=[1, 2, 3],
timestamps={"a": [1, 2, 3, 4, 5]},
seq_num=[1, 2, 3, 4, 5],
)


Expand Down

0 comments on commit 2910dfe

Please sign in to comment.