Skip to content

Commit

Permalink
Give _config an initial value, add a copy method
Browse files Browse the repository at this point in the history
  • Loading branch information
dagardner-nv committed Aug 20, 2024
1 parent 13681ea commit aa13a0e
Showing 1 changed file with 20 additions and 0 deletions.
20 changes: 20 additions & 0 deletions python/morpheus/morpheus/messages/control_message.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@ class ControlMessage(MessageBase, cpp_class=_messages.ControlMessage):
def __init__(self, config: dict = None):
super().__init__()

self._config: dict = {"metadata": {}}

self._payload: MessageMeta = None
self._tensors: TensorMemory = None

Expand Down Expand Up @@ -72,6 +74,24 @@ def config(self, config: dict = None) -> dict:

return self._config

def copy(self) -> "ControlMessage":
config = self._config.copy()
config["type"] = self.task_type().name

tasks = []
for (task_type, task_queue) in self.get_tasks().items():
for task in task_queue:
tasks.append({"type": task_type, "properties": task})

config["tasks"] = tasks

new_cm = ControlMessage(config)
new_cm._payload = self._payload
new_cm._tensors = self._tensors
new_cm._timestamps = self._timestamps.copy()

return new_cm

def has_task(self, task_type: str) -> bool:
"""
Return True if the control message has at least one task of the given type
Expand Down

0 comments on commit aa13a0e

Please sign in to comment.