def reconstruct_artifact_multimap()

in tfx/orchestration/portable/mlmd/event_lib.py [0:0]

• 25 lines of code
• 12 McCabe index (conditional complexity)

def reconstruct_artifact_multimap(
    artifacts: Sequence[_Artifact],
    events: Sequence[metadata_store_pb2.Event]) -> _ArtifactMultiDict:
  """Reconstructs input or output artifact maps from events."""
  execution_ids = {e.execution_id for e in events}
  events_by_artifact_id = {e.artifact_id: e for e in events}
  if len(execution_ids) > 1:
    raise ValueError(
        'All events should be from the same execution but got: '
        f'{execution_ids}.')

  artifacts_by_id = {a.id: a for a in artifacts}
  artifact_id_multimap = reconstruct_artifact_id_multimap(events)
  result = {
      key: [artifacts_by_id[i] for i in artifact_ids]
      for key, artifact_ids in artifact_id_multimap.items()
  }
  for key, artifacts in result.items():
    artifact_types = {a.type_id for a in artifacts}
    if len(artifact_types) != 1:
      raise ValueError(
          f'Artifact type of key "{key}" is heterogeneous: {artifact_types}')
    event_types = {events_by_artifact_id[a.id].type for a in artifacts}
    if len(event_types) != 1:
      raise ValueError(
          f'Event type of key "{key}" is heterogeneous: {event_types}')
  return result