chore: import upstream snapshot with attribution

python/ray/dag/format_utils.py

@@ -0,0 +1,155 @@
+from ray.dag import DAGNode
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+def get_dag_node_str(
+    dag_node: DAGNode,
+    body_line,
+):
+    indent = _get_indentation()
+    other_args_to_resolve_lines = _get_other_args_to_resolve_lines(
+        dag_node._bound_other_args_to_resolve
+    )
+    return (
+        f"({dag_node.__class__.__name__}, {dag_node._stable_uuid})(\n"
+        f"{indent}body={body_line}\n"
+        f"{indent}args={_get_args_lines(dag_node._bound_args)}\n"
+        f"{indent}kwargs={_get_kwargs_lines(dag_node._bound_kwargs)}\n"
+        f"{indent}options={_get_options_lines(dag_node._bound_options)}\n"
+        f"{indent}other_args_to_resolve={other_args_to_resolve_lines}\n"
+        f")"
+    )
+
+
+def _get_indentation(num_spaces=4):
+    return " " * num_spaces
+
+
+def _get_args_lines(bound_args):
+    """Pretty prints bounded args of a DAGNode, and recursively handle
+    DAGNode in list / dict containers.
+    """
+    indent = _get_indentation()
+    lines = []
+    for arg in bound_args:
+        if isinstance(arg, DAGNode):
+            node_repr_lines = str(arg).split("\n")
+            for node_repr_line in node_repr_lines:
+                lines.append(f"{indent}" + node_repr_line)
+        elif isinstance(arg, list):
+            for ele in arg:
+                node_repr_lines = str(ele).split("\n")
+                for node_repr_line in node_repr_lines:
+                    lines.append(f"{indent}" + node_repr_line)
+        elif isinstance(arg, dict):
+            for _, val in arg.items():
+                node_repr_lines = str(val).split("\n")
+                for node_repr_line in node_repr_lines:
+                    lines.append(f"{indent}" + node_repr_line)
+        # TODO: (jiaodong) Handle nested containers and other obj types
+        else:
+            lines.append(f"{indent}" + str(arg) + ", ")
+
+    if len(lines) == 0:
+        args_line = "[]"
+    else:
+        args_line = "["
+        for args in lines:
+            args_line += f"\n{indent}{args}"
+        args_line += f"\n{indent}]"
+
+    return args_line
+
+
+def _get_kwargs_lines(bound_kwargs):
+    """Pretty prints bounded kwargs of a DAGNode, and recursively handle
+    DAGNode in list / dict containers.
+    """
+    # TODO: (jiaodong) Nits, we're missing keys and indentation was a bit off.
+    if not bound_kwargs:
+        return "{}"
+    indent = _get_indentation()
+    kwargs_lines = []
+    for key, val in bound_kwargs.items():
+        if isinstance(val, DAGNode):
+            node_repr_lines = str(val).split("\n")
+            for index, node_repr_line in enumerate(node_repr_lines):
+                if index == 0:
+                    kwargs_lines.append(
+                        f"{indent}{key}:" + f"{indent}" + node_repr_line
+                    )
+                else:
+                    kwargs_lines.append(f"{indent}{indent}" + node_repr_line)
+
+        elif isinstance(val, list):
+            for ele in val:
+                node_repr_lines = str(ele).split("\n")
+                for node_repr_line in node_repr_lines:
+                    kwargs_lines.append(f"{indent}" + node_repr_line)
+        elif isinstance(val, dict):
+            for _, inner_val in val.items():
+                node_repr_lines = str(inner_val).split("\n")
+                for node_repr_line in node_repr_lines:
+                    kwargs_lines.append(f"{indent}" + node_repr_line)
+        # TODO: (jiaodong) Handle nested containers and other obj types
+        else:
+            kwargs_lines.append(val)
+
+    if len(kwargs_lines) > 0:
+        kwargs_line = "{"
+        for line in kwargs_lines:
+            kwargs_line += f"\n{indent}{line}"
+        kwargs_line += f"\n{indent}}}"
+    else:
+        kwargs_line = "{}"
+
+    return kwargs_line
+
+
+def _get_options_lines(bound_options):
+    """Pretty prints .options() in DAGNode. Only prints non-empty values."""
+    if not bound_options:
+        return "{}"
+    indent = _get_indentation()
+    options_lines = []
+    for key, val in bound_options.items():
+        if val:
+            options_lines.append(f"{indent}{key}: " + str(val))
+
+    options_line = "{"
+    for line in options_lines:
+        options_line += f"\n{indent}{line}"
+    options_line += f"\n{indent}}}"
+    return options_line
+
+
+def _get_other_args_to_resolve_lines(other_args_to_resolve):
+    if not other_args_to_resolve:
+        return "{}"
+    indent = _get_indentation()
+    other_args_to_resolve_lines = []
+    for key, val in other_args_to_resolve.items():
+        if isinstance(val, DAGNode):
+            node_repr_lines = str(val).split("\n")
+            for index, node_repr_line in enumerate(node_repr_lines):
+                if index == 0:
+                    other_args_to_resolve_lines.append(
+                        f"{indent}{key}:"
+                        + f"{indent}"
+                        + "\n"
+                        + f"{indent}{indent}{indent}"
+                        + node_repr_line
+                    )
+                else:
+                    other_args_to_resolve_lines.append(
+                        f"{indent}{indent}" + node_repr_line
+                    )
+        else:
+            other_args_to_resolve_lines.append(f"{indent}{key}: " + str(val))
+
+    other_args_to_resolve_line = "{"
+    for line in other_args_to_resolve_lines:
+        other_args_to_resolve_line += f"\n{indent}{line}"
+    other_args_to_resolve_line += f"\n{indent}}}"
+    return other_args_to_resolve_line