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

Author: stainless-app[bot]

src/kernel/_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/kernel/_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/kernel/_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)

src/kernel/resources/browsers/browsers.py

@@ -49,8 +49,9 @@
     ReplaysResourceWithStreamingResponse,
     AsyncReplaysResourceWithStreamingResponse,
 )
+from ..._files import deepcopy_with_paths
 from ..._types import Body, Omit, Query, Headers, NoneType, NotGiven, omit, not_given
-from ..._utils import extract_files, path_template, maybe_transform, deepcopy_minimal, async_maybe_transform
+from ..._utils import extract_files, path_template, maybe_transform, async_maybe_transform
 from .computer import (
     ComputerResource,
     AsyncComputerResource,
@@ -573,7 +574,7 @@ def load_extensions(
         if not id:
             raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
         extra_headers = {"Accept": "*/*", **(extra_headers or {})}
-        body = deepcopy_minimal({"extensions": extensions})
+        body = deepcopy_with_paths({"extensions": extensions}, [["extensions", "<array>", "zip_file"]])
         files = extract_files(cast(Mapping[str, object], body), paths=[["extensions", "<array>", "zip_file"]])
         # It should be noted that the actual Content-Type header that will be
         # sent to the server will contain a `boundary` parameter, e.g.
@@ -1075,7 +1076,7 @@ async def load_extensions(
         if not id:
             raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
         extra_headers = {"Accept": "*/*", **(extra_headers or {})}
-        body = deepcopy_minimal({"extensions": extensions})
+        body = deepcopy_with_paths({"extensions": extensions}, [["extensions", "<array>", "zip_file"]])
         files = extract_files(cast(Mapping[str, object], body), paths=[["extensions", "<array>", "zip_file"]])
         # It should be noted that the actual Content-Type header that will be
         # sent to the server will contain a `boundary` parameter, e.g.

src/kernel/resources/browsers/fs/fs.py

@@ -15,7 +15,7 @@
     WatchResourceWithStreamingResponse,
     AsyncWatchResourceWithStreamingResponse,
 )
-from ...._files import read_file_content, async_read_file_content
+from ...._files import read_file_content, deepcopy_with_paths, async_read_file_content
 from ...._types import (
     Body,
     Omit,
@@ -30,7 +30,7 @@
     omit,
     not_given,
 )
-from ...._utils import extract_files, path_template, maybe_transform, deepcopy_minimal, async_maybe_transform
+from ...._utils import extract_files, path_template, maybe_transform, async_maybe_transform
 from ...._compat import cached_property
 from ...._resource import SyncAPIResource, AsyncAPIResource
 from ...._response import (
@@ -509,7 +509,7 @@ def upload(
         if not id:
             raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
         extra_headers = {"Accept": "*/*", **(extra_headers or {})}
-        body = deepcopy_minimal({"files": files})
+        body = deepcopy_with_paths({"files": files}, [["files", "<array>", "file"]])
         extracted_files = extract_files(cast(Mapping[str, object], body), paths=[["files", "<array>", "file"]])
         # It should be noted that the actual Content-Type header that will be
         # sent to the server will contain a `boundary` parameter, e.g.
@@ -555,11 +555,12 @@ def upload_zip(
         if not id:
             raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
         extra_headers = {"Accept": "*/*", **(extra_headers or {})}
-        body = deepcopy_minimal(
+        body = deepcopy_with_paths(
             {
                 "dest_path": dest_path,
                 "zip_file": zip_file,
-            }
+            },
+            [["zip_file"]],
         )
         files = extract_files(cast(Mapping[str, object], body), paths=[["zip_file"]])
         # It should be noted that the actual Content-Type header that will be
@@ -1071,7 +1072,7 @@ async def upload(
         if not id:
             raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
         extra_headers = {"Accept": "*/*", **(extra_headers or {})}
-        body = deepcopy_minimal({"files": files})
+        body = deepcopy_with_paths({"files": files}, [["files", "<array>", "file"]])
         extracted_files = extract_files(cast(Mapping[str, object], body), paths=[["files", "<array>", "file"]])
         # It should be noted that the actual Content-Type header that will be
         # sent to the server will contain a `boundary` parameter, e.g.
@@ -1117,11 +1118,12 @@ async def upload_zip(
         if not id:
             raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
         extra_headers = {"Accept": "*/*", **(extra_headers or {})}
-        body = deepcopy_minimal(
+        body = deepcopy_with_paths(
             {
                 "dest_path": dest_path,
                 "zip_file": zip_file,
-            }
+            },
+            [["zip_file"]],
         )
         files = extract_files(cast(Mapping[str, object], body), paths=[["zip_file"]])
         # It should be noted that the actual Content-Type header that will be

src/kernel/resources/deployments.py

@@ -8,8 +8,9 @@
 import httpx
 
 from ..types import deployment_list_params, deployment_create_params, deployment_follow_params
+from .._files import deepcopy_with_paths
 from .._types import Body, Omit, Query, Headers, NoneType, NotGiven, FileTypes, omit, not_given
-from .._utils import extract_files, path_template, maybe_transform, deepcopy_minimal, async_maybe_transform
+from .._utils import extract_files, path_template, maybe_transform, async_maybe_transform
 from .._compat import cached_property
 from .._resource import SyncAPIResource, AsyncAPIResource
 from .._response import (
@@ -95,7 +96,7 @@ def create(
 
           timeout: Override the client-level default timeout for this request, in seconds
         """
-        body = deepcopy_minimal(
+        body = deepcopy_with_paths(
             {
                 "entrypoint_rel_path": entrypoint_rel_path,
                 "env_vars": env_vars,
@@ -104,7 +105,8 @@ def create(
                 "region": region,
                 "source": source,
                 "version": version,
-            }
+            },
+            [["file"]],
         )
         files = extract_files(cast(Mapping[str, object], body), paths=[["file"]])
         # It should be noted that the actual Content-Type header that will be
@@ -360,7 +362,7 @@ async def create(
 
           timeout: Override the client-level default timeout for this request, in seconds
         """
-        body = deepcopy_minimal(
+        body = deepcopy_with_paths(
             {
                 "entrypoint_rel_path": entrypoint_rel_path,
                 "env_vars": env_vars,
@@ -369,7 +371,8 @@ async def create(
                 "region": region,
                 "source": source,
                 "version": version,
-            }
+            },
+            [["file"]],
         )
         files = extract_files(cast(Mapping[str, object], body), paths=[["file"]])
         # It should be noted that the actual Content-Type header that will be

src/kernel/resources/extensions.py

@@ -8,8 +8,9 @@
 import httpx
 
 from ..types import extension_upload_params, extension_download_from_chrome_store_params
+from .._files import deepcopy_with_paths
 from .._types import Body, Omit, Query, Headers, NoneType, NotGiven, FileTypes, omit, not_given
-from .._utils import extract_files, path_template, maybe_transform, deepcopy_minimal, async_maybe_transform
+from .._utils import extract_files, path_template, maybe_transform, async_maybe_transform
 from .._compat import cached_property
 from .._resource import SyncAPIResource, AsyncAPIResource
 from .._response import (
@@ -220,11 +221,12 @@ def upload(
 
           timeout: Override the client-level default timeout for this request, in seconds
         """
-        body = deepcopy_minimal(
+        body = deepcopy_with_paths(
             {
                 "file": file,
                 "name": name,
-            }
+            },
+            [["file"]],
         )
         files = extract_files(cast(Mapping[str, object], body), paths=[["file"]])
         # It should be noted that the actual Content-Type header that will be
@@ -429,11 +431,12 @@ async def upload(
 
           timeout: Override the client-level default timeout for this request, in seconds
         """
-        body = deepcopy_minimal(
+        body = deepcopy_with_paths(
             {
                 "file": file,
                 "name": name,
-            }
+            },
+            [["file"]],
         )
         files = extract_files(cast(Mapping[str, object], body), paths=[["file"]])
         # It should be noted that the actual Content-Type header that will be