houqp commented on a change in pull request #10153:
URL: https://github.com/apache/airflow/pull/10153#discussion_r471805675
##########
File path: airflow/serialization/serialized_objects.py
##########
@@ -626,3 +653,60 @@ 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) -> Union[Dict[str,
Any], None]:
+ """Serializes TaskGroup into a JSON object.
+ """
+ if not task_group:
+ return None
+
+ serialize_group = {}
+ serialize_group["_group_id"] = task_group._group_id # pylint:
disable=protected-access
+
+ serialize_group['children'] = { # type: ignore
+ 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()
+ }
+ serialize_group['tooltip'] = task_group.tooltip
+ serialize_group['ui_color'] = task_group.ui_color
+ serialize_group['ui_fgcolor'] = task_group.ui_fgcolor
+ return serialize_group
+
+ @classmethod
+ def deserialize_task_group(
+ cls,
+ encoded_group: Dict[str, Any],
+ parent_group: Union[TaskGroup, None],
+ task_dict: Dict[str, BaseOperator]
+ ) -> Union[TaskGroup, None]:
Review comment:
For nullable types, it's better to use `Optional`
```suggestion
) -> Optional[TaskGroup]:
```
##########
File path: airflow/models/baseoperator.py
##########
@@ -359,9 +363,12 @@ def __init__(
do_xcom_push: bool = True,
inlets: Optional[Any] = None,
outlets: Optional[Any] = None,
+ task_group: Optional["TaskGroup"] = None,
**kwargs
):
from airflow.models.dag import DagContext
+ from airflow.utils.task_group import TaskGroupContext
Review comment:
is this for avoiding circular import? I thought having the `if
TYPE_CHECKING` in utils.task_group should be enough?
##########
File path: airflow/models/baseoperator.py
##########
@@ -380,7 +387,15 @@ def __init__(
stacklevel=3
)
validate_key(task_id)
- self.task_id = task_id
+ self._task_id = task_id
Review comment:
I am a bit concerned with us changing the semantic of `self.task_id`
here for a UI feature. There are airflow code written with the assumption this
is the id that can be used for querying in database.
If the task_group prefixed "task_id" is only intended for visualization,
would it make more sense to create a new attribute for it instead of overriding
the existing one?
##########
File path: airflow/serialization/serialized_objects.py
##########
@@ -626,3 +653,60 @@ 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) -> Union[Dict[str,
Any], None]:
+ """Serializes TaskGroup into a JSON object.
+ """
+ if not task_group:
+ return None
+
+ serialize_group = {}
+ serialize_group["_group_id"] = task_group._group_id # pylint:
disable=protected-access
+
+ serialize_group['children'] = { # type: ignore
+ 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()
+ }
+ serialize_group['tooltip'] = task_group.tooltip
+ serialize_group['ui_color'] = task_group.ui_color
+ serialize_group['ui_fgcolor'] = task_group.ui_fgcolor
Review comment:
nitpick, looks like these individual dictionary set operation can be
combined into the dict initialization statement?
```python
serialize_group = {
"tooltip": task_group.toolip,
...
}
```
##########
File path: airflow/www/views.py
##########
@@ -147,6 +147,41 @@ def get_date_time_num_runs_dag_runs_form_data(request,
session, dag):
}
+def task_group_to_dict(task_group):
+ """
+ Create a nested dict representation of this TaskGroup and its children
used for
+ rendering web UI.
+ """
+ from airflow.models.baseoperator import BaseOperator
+
+ if isinstance(task_group, BaseOperator):
+ return {
+ 'id': task_group.task_id,
+ 'value': {
+ 'label': task_group.label,
+ 'labelStyle': "fill:{0};".format(task_group.ui_fgcolor),
+ 'style': "fill:{0};".format(task_group.ui_color),
+ 'rx': 5,
+ 'ry': 5,
+ }
+ }
+
+ return {
+ "id": task_group.group_id,
+ 'value': {
+ 'label': task_group.label,
+ 'labelStyle': "fill:{0};".format(task_group.ui_fgcolor),
Review comment:
nitpick, would be good to consistently use `f-string` everywhere instead
of format, I noticed you are already using f-string for style key a line below.
----------------------------------------------------------------
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]
