perf(client): optimize file structure copying in multipart requests

Author: stainless-app[bot]

src/gitpod/_files.py

@@ -3,8 +3,8 @@
 import io
 import os
 import pathlib
-from typing import overload
-from typing_extensions import TypeGuard
+from typing import Sequence, cast, overload
+from typing_extensions import TypeVar, TypeGuard
 
 import anyio
 
@@ -17,7 +17,9 @@
     HttpxFileContent,
     HttpxRequestFiles,
 )
-from ._utils import is_tuple_t, is_mapping_t, is_sequence_t
+from ._utils import is_list, is_mapping, is_tuple_t, is_mapping_t, is_sequence_t
+
+_T = TypeVar("_T")
 
 
 def is_base64_file_input(obj: object) -> TypeGuard[Base64FileInput]:
@@ -121,3 +123,51 @@ async def async_read_file_content(file: FileContent) -> HttpxFileContent:
         return await anyio.Path(file).read_bytes()
 
     return file
+
+
+def deepcopy_with_paths(item: _T, paths: Sequence[Sequence[str]]) -> _T:
+    """Copy only the containers along the given paths.
+
+    Used to guard against mutation by extract_files without copying the entire structure.
+    Only dicts and lists that lie on a path are copied; everything else
+    is returned by reference.
+
+    For example, given paths=[["foo", "files", "file"]] and the structure:
+        {
+            "foo": {
+                "bar": {"baz": {}},
+                "files": {"file": <content>}
+            }
+        }
+    The root dict, "foo", and "files" are copied (they lie on the path).
+    "bar" and "baz" are returned by reference (off the path).
+    """
+    return _deepcopy_with_paths(item, paths, 0)
+
+
+def _deepcopy_with_paths(item: _T, paths: Sequence[Sequence[str]], index: int) -> _T:
+    if not paths:
+        return item
+    if is_mapping(item):
+        key_to_paths: dict[str, list[Sequence[str]]] = {}
+        for path in paths:
+            if index < len(path):
+                key_to_paths.setdefault(path[index], []).append(path)
+
+        # if no path continues through this mapping, it won't be mutated and copying it is redundant
+        if not key_to_paths:
+            return item
+
+        result = dict(item)
+        for key, subpaths in key_to_paths.items():
+            if key in result:
+                result[key] = _deepcopy_with_paths(result[key], subpaths, index + 1)
+        return cast(_T, result)
+    if is_list(item):
+        array_paths = [path for path in paths if index < len(path) and path[index] == "<array>"]
+
+        # if no path expects a list here, nothing will be mutated inside it - return by reference
+        if not array_paths:
+            return cast(_T, item)
+        return cast(_T, [_deepcopy_with_paths(entry, array_paths, index + 1) for entry in item])
+    return item

src/gitpod/_utils/__init__.py

@@ -24,7 +24,6 @@
     coerce_integer as coerce_integer,
     file_from_path as file_from_path,
     strip_not_given as strip_not_given,
-    deepcopy_minimal as deepcopy_minimal,
     get_async_library as get_async_library,
     maybe_coerce_float as maybe_coerce_float,
     get_required_header as get_required_header,

src/gitpod/_utils/_utils.py

@@ -177,21 +177,6 @@ def is_iterable(obj: object) -> TypeGuard[Iterable[object]]:
     return isinstance(obj, Iterable)
 
 
-def deepcopy_minimal(item: _T) -> _T:
-    """Minimal reimplementation of copy.deepcopy() that will only copy certain object types:
-
-    - mappings, e.g. `dict`
-    - list
-
-    This is done for performance reasons.
-    """
-    if is_mapping(item):
-        return cast(_T, {k: deepcopy_minimal(v) for k, v in item.items()})
-    if is_list(item):
-        return cast(_T, [deepcopy_minimal(entry) for entry in item])
-    return item
-
-
 # copied from https://github.com/Rapptz/RoboDanny
 def human_join(seq: Sequence[str], *, delim: str = ", ", final: str = "or") -> str:
     size = len(seq)

tests/test_deepcopy.py

@@ -4,7 +4,8 @@
 import pytest
 from dirty_equals import IsDict, IsList, IsBytes, IsTuple
 
-from gitpod._files import to_httpx_files, async_to_httpx_files
+from gitpod._files import to_httpx_files, deepcopy_with_paths, async_to_httpx_files
+from gitpod._utils import extract_files
 
 readme_path = Path(__file__).parent.parent.joinpath("README.md")
 
@@ -49,3 +50,99 @@ def test_string_not_allowed() -> None:
                 "file": "foo",  # type: ignore
             }
         )
+
+
+def assert_different_identities(obj1: object, obj2: object) -> None:
+    assert obj1 == obj2
+    assert obj1 is not obj2
+
+
+class TestDeepcopyWithPaths:
+    def test_copies_top_level_dict(self) -> None:
+        original = {"file": b"data", "other": "value"}
+        result = deepcopy_with_paths(original, [["file"]])
+        assert_different_identities(result, original)
+
+    def test_file_value_is_same_reference(self) -> None:
+        file_bytes = b"contents"
+        original = {"file": file_bytes}
+        result = deepcopy_with_paths(original, [["file"]])
+        assert_different_identities(result, original)
+        assert result["file"] is file_bytes
+
+    def test_list_popped_wholesale(self) -> None:
+        files = [b"f1", b"f2"]
+        original = {"files": files, "title": "t"}
+        result = deepcopy_with_paths(original, [["files", "<array>"]])
+        assert_different_identities(result, original)
+        result_files = result["files"]
+        assert isinstance(result_files, list)
+        assert_different_identities(result_files, files)
+
+    def test_nested_array_path_copies_list_and_elements(self) -> None:
+        elem1 = {"file": b"f1", "extra": 1}
+        elem2 = {"file": b"f2", "extra": 2}
+        original = {"items": [elem1, elem2]}
+        result = deepcopy_with_paths(original, [["items", "<array>", "file"]])
+        assert_different_identities(result, original)
+        result_items = result["items"]
+        assert isinstance(result_items, list)
+        assert_different_identities(result_items, original["items"])
+        assert_different_identities(result_items[0], elem1)
+        assert_different_identities(result_items[1], elem2)
+
+    def test_empty_paths_returns_same_object(self) -> None:
+        original = {"foo": "bar"}
+        result = deepcopy_with_paths(original, [])
+        assert result is original
+
+    def test_multiple_paths(self) -> None:
+        f1 = b"file1"
+        f2 = b"file2"
+        original = {"a": f1, "b": f2, "c": "unchanged"}
+        result = deepcopy_with_paths(original, [["a"], ["b"]])
+        assert_different_identities(result, original)
+        assert result["a"] is f1
+        assert result["b"] is f2
+        assert result["c"] is original["c"]
+
+    def test_extract_files_does_not_mutate_original_top_level(self) -> None:
+        file_bytes = b"contents"
+        original = {"file": file_bytes, "other": "value"}
+
+        copied = deepcopy_with_paths(original, [["file"]])
+        extracted = extract_files(copied, paths=[["file"]])
+
+        assert extracted == [("file", file_bytes)]
+        assert original == {"file": file_bytes, "other": "value"}
+        assert copied == {"other": "value"}
+
+    def test_extract_files_does_not_mutate_original_nested_array_path(self) -> None:
+        file1 = b"f1"
+        file2 = b"f2"
+        original = {
+            "items": [
+                {"file": file1, "extra": 1},
+                {"file": file2, "extra": 2},
+            ],
+            "title": "example",
+        }
+
+        copied = deepcopy_with_paths(original, [["items", "<array>", "file"]])
+        extracted = extract_files(copied, paths=[["items", "<array>", "file"]])
+
+        assert extracted == [("items[][file]", file1), ("items[][file]", file2)]
+        assert original == {
+            "items": [
+                {"file": file1, "extra": 1},
+                {"file": file2, "extra": 2},
+            ],
+            "title": "example",
+        }
+        assert copied == {
+            "items": [
+                {"extra": 1},
+                {"extra": 2},
+            ],
+            "title": "example",
+        }