turbaszek commented on a change in pull request #10153:
URL: https://github.com/apache/airflow/pull/10153#discussion_r484262687
##########
File path: airflow/serialization/serialized_objects.py
##########
@@ -625,3 +647,69 @@ def from_dict(cls, serialized_obj: dict) ->
'SerializedDAG':
if ver != cls.SERIALIZER_VERSION:
raise ValueError("Unsure how to deserialize version
{!r}".format(ver))
return cls.deserialize_dag(serialized_obj['dag'])
+
+
+class SerializedTaskGroup(TaskGroup, BaseSerialization):
+ """
+ A JSON serializable representation of TaskGroup.
+ """
+ @classmethod
+ def serialize_task_group(cls, task_group: TaskGroup) ->
Optional[Union[Dict[str, Any]]]:
+ """Serializes TaskGroup into a JSON object.
+ """
+ if not task_group:
+ return None
+
+ serialize_group = {
+ "_group_id": task_group._group_id, # pylint:
disable=protected-access
+ "prefix_group_id": task_group.prefix_group_id,
+ "tooltip": task_group.tooltip,
+ "ui_color": task_group.ui_color,
+ "ui_fgcolor": task_group.ui_fgcolor,
+ "children": {
+ label: (DAT.OP, child.task_id)
+ if isinstance(child, BaseOperator) else
+ (DAT.TASK_GROUP,
SerializedTaskGroup.serialize_task_group(child))
+ for label, child in task_group.children.items()
+ },
+ "upstream_group_ids":
cls._serialize(list(task_group.upstream_group_ids)),
+ "downstream_group_ids":
cls._serialize(list(task_group.downstream_group_ids)),
+ "upstream_task_ids":
cls._serialize(list(task_group.upstream_task_ids)),
+ "downstream_task_ids":
cls._serialize(list(task_group.downstream_task_ids)),
+
+ }
+
+ return serialize_group
+
+ @classmethod
+ def deserialize_task_group(
+ cls,
+ encoded_group: Dict[str, Any],
+ parent_group: Optional[TaskGroup],
+ task_dict: Dict[str, BaseOperator]
+ ) -> Optional[TaskGroup]:
+ """Deserializes a TaskGroup from a JSON object.
Review comment:
```suggestion
"""
Deserializes a TaskGroup from a JSON object.
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
