From f6bd78ae62f579d0be5621ffbff7d17a5595919d Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Fri, 7 Jun 2024 16:09:03 -0400
Subject: [PATCH] feat: port `util.copy_task` from `gecko_taskgraph`

And use it in `util.templates.merge`.

This is a port of the following patch in `gecko_taskgraph`:
https://hg.mozilla.org/mozilla-central/rev/fa2cdac6989ac78d606d8f8adb10d826a6e50429
---
 src/taskgraph/util/copy.py      | 47 +++++++++++++++++++++++++++++++++
 src/taskgraph/util/templates.py |  4 +--
 test/test_util_copy.py          | 34 ++++++++++++++++++++++++
 3 files changed, 83 insertions(+), 2 deletions(-)
 create mode 100644 src/taskgraph/util/copy.py
 create mode 100644 test/test_util_copy.py

diff --git a/src/taskgraph/util/copy.py b/src/taskgraph/util/copy.py
new file mode 100644
index 000000000..a8d99a24b
--- /dev/null
+++ b/src/taskgraph/util/copy.py
@@ -0,0 +1,47 @@
+from typing import Any
+
+from taskgraph.task import Task
+from taskgraph.util.readonlydict import ReadOnlyDict
+
+immutable_types = {int, float, bool, str, type(None), ReadOnlyDict}
+
+
+def deepcopy(obj: Any) -> Any:
+    """Perform a deep copy of an object with a tree like structure.
+
+    This is a re-implementation of Python's `copy.deepcopy` function with a few key differences:
+
+    1. Unlike the stdlib, this does *not* support copying graph-like structure,
+    which allows it to be more efficient than deepcopy on tree-like structures
+    (such as Tasks).
+    2. This special cases support for `taskgraph.task.Task` objects.
+
+    Args:
+        obj: The object to deep copy.
+
+    Returns:
+        A deep copy of the object.
+    """
+    ty = type(obj)
+    if ty in immutable_types:
+        return obj
+    if ty is dict:
+        return {k: deepcopy(v) for k, v in obj.items()}
+    if ty is list:
+        return [deepcopy(elt) for elt in obj]
+    if ty is Task:
+        task = Task(
+            kind=deepcopy(obj.kind),
+            label=deepcopy(obj.label),
+            attributes=deepcopy(obj.attributes),
+            task=deepcopy(obj.task),
+            description=deepcopy(obj.description),
+            optimization=deepcopy(obj.optimization),
+            dependencies=deepcopy(obj.dependencies),
+            soft_dependencies=deepcopy(obj.soft_dependencies),
+            if_dependencies=deepcopy(obj.if_dependencies),
+        )
+        if obj.task_id:
+            task.task_id = obj.task_id
+        return task
+    raise NotImplementedError(f"copying '{ty}' from '{obj}'")
diff --git a/src/taskgraph/util/templates.py b/src/taskgraph/util/templates.py
index d15459dc5..be5fc4cbc 100644
--- a/src/taskgraph/util/templates.py
+++ b/src/taskgraph/util/templates.py
@@ -2,7 +2,7 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-import copy
+from taskgraph.util.copy import deepcopy
 
 
 def merge_to(source, dest):
@@ -55,7 +55,7 @@ def merge(*objects):
     Returns the result without modifying any arguments.
     """
     if len(objects) == 1:
-        return copy.deepcopy(objects[0])
+        return deepcopy(objects[0])
     return merge_to(objects[-1], merge(*objects[:-1]))
 
 
diff --git a/test/test_util_copy.py b/test/test_util_copy.py
new file mode 100644
index 000000000..d541f035c
--- /dev/null
+++ b/test/test_util_copy.py
@@ -0,0 +1,34 @@
+import pytest
+
+from taskgraph.task import Task
+from taskgraph.util.copy import deepcopy, immutable_types
+from taskgraph.util.readonlydict import ReadOnlyDict
+
+
+@pytest.mark.parametrize(
+    "input",
+    (
+        1,
+        False,
+        "foo",
+        ReadOnlyDict(a=1, b="foo"),
+        ["foo", "bar"],
+        {
+            "foo": Task(
+                label="abc",
+                kind="kind",
+                attributes={"bar": "baz"},
+                dependencies={"dep": "bar"},
+                task={"payload": {"command": ["echo hello"]}},
+            )
+        },
+    ),
+)
+def test_deepcopy(input):
+    result = deepcopy(input)
+    assert result == input
+
+    if type(result) in immutable_types:
+        assert id(result) == id(input)
+    else:
+        assert id(result) != id(input)