fix(core): Refactor copy file by Tranquility2 · Pull Request #996 · testcontainers/testcontainers-python · GitHub
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions core/README.rst
63 changes: 15 additions & 48 deletions core/testcontainers/core/container.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from testcontainers.core.exceptions import ContainerConnectException, ContainerStartException
from testcontainers.core.labels import LABEL_SESSION_ID, SESSION_ID
from testcontainers.core.network import Network
from testcontainers.core.transferable import Transferable, TransferSpec
from testcontainers.core.transferable import Transferable, TransferSpec, build_transfer_tar
from testcontainers.core.utils import is_arm, setup_logger
from testcontainers.core.wait_strategies import LogMessageWaitStrategy
from testcontainers.core.waiting_utils import WaitStrategy
Expand Down Expand Up @@ -289,7 +289,7 @@ def with_volume_mapping(self, host: Union[str, PathLike[str]], container: str, m

def with_tmpfs_mount(self, container_path: str, size: Optional[str] = None) -> Self:
"""Mount a tmpfs volume on the container.

:param container_path: Container path to mount tmpfs on (e.g., '/data')
:param size: Optional size limit (e.g., '256m', '1g'). If None, unbounded.
:return: Self for chaining
Expand Down Expand Up @@ -342,57 +342,24 @@ def copy_into_container(self, transferable: Transferable, destination_in_contain
return self._transfer_into_container(transferable, destination_in_container, mode)

def _transfer_into_container(self, transferable: Transferable, destination_in_container: str, mode: int) -> None:
if isinstance(transferable, bytes):
self._transfer_file_content_into_container(transferable, destination_in_container, mode)
elif isinstance(transferable, pathlib.Path):
if transferable.is_file():
self._transfer_file_content_into_container(transferable.read_bytes(), destination_in_container, mode)
elif transferable.is_dir():
self._transfer_directory_into_container(transferable, destination_in_container, mode)
else:
raise TypeError(f"Path {transferable} is neither a file nor directory")
else:
raise TypeError("source must be bytes or PathLike")

def _transfer_file_content_into_container(
self, file_content: bytes, destination_in_container: str, mode: int
) -> None:
fileobj = io.BytesIO()
with tarfile.open(fileobj=fileobj, mode="w") as tar:
tarinfo = tarfile.TarInfo(name=destination_in_container)
tarinfo.size = len(file_content)
tarinfo.mode = mode
tar.addfile(tarinfo, io.BytesIO(file_content))
fileobj.seek(0)
assert self._container is not None
rv = self._container.put_archive(path="/", data=fileobj.getvalue())
assert rv is True

def _transfer_directory_into_container(
self, source_directory: pathlib.Path, destination_in_container: str, mode: int
) -> None:
assert self._container is not None
result = self._container.exec_run(["mkdir", "-p", destination_in_container])
assert result.exit_code == 0
if not self._container:
raise ContainerStartException("Container must be started before transferring files")

fileobj = io.BytesIO()
with tarfile.open(fileobj=fileobj, mode="w") as tar:
tar.add(source_directory, arcname=source_directory.name)
fileobj.seek(0)
rv = self._container.put_archive(path=destination_in_container, data=fileobj.getvalue())
assert rv is True
data = build_transfer_tar(transferable, destination_in_container, mode)
if not self._container.put_archive(path="/", data=data):
raise OSError(f"Failed to put archive into container at {destination_in_container}")

def copy_from_container(self, source_in_container: str, destination_on_host: pathlib.Path) -> None:
assert self._container is not None
if not self._container:
raise ContainerStartException("Container must be started before copying files")

tar_stream, _ = self._container.get_archive(source_in_container)

for chunk in tar_stream:
with tarfile.open(fileobj=io.BytesIO(chunk)) as tar:
for member in tar.getmembers():
with open(destination_on_host, "wb") as f:
fileobj = tar.extractfile(member)
assert fileobj is not None
f.write(fileobj.read())
with tarfile.open(fileobj=io.BytesIO(b"".join(tar_stream))) as tar:
for member in tar.getmembers():
extracted = tar.extractfile(member)
if extracted is not None:
destination_on_host.write_bytes(extracted.read())


class Reaper:
Expand Down
27 changes: 27 additions & 0 deletions core/testcontainers/core/transferable.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,33 @@
import io
import pathlib
import tarfile
from typing import Union

Transferable = Union[bytes, pathlib.Path]

TransferSpec = Union[tuple[Transferable, str], tuple[Transferable, str, int]]


def build_transfer_tar(transferable: Transferable, destination: str, mode: int = 0o644) -> bytes:
"""Build a tar archive containing the transferable, ready for put_archive(path="/")."""
buf = io.BytesIO()
with tarfile.open(fileobj=buf, mode="w") as tar:
if isinstance(transferable, bytes):
info = tarfile.TarInfo(name=destination)
info.size = len(transferable)
info.mode = mode
tar.addfile(info, io.BytesIO(transferable))
elif isinstance(transferable, pathlib.Path):
if transferable.is_file():
info = tarfile.TarInfo(name=destination)
info.size = transferable.stat().st_size
info.mode = mode
with transferable.open("rb") as f:
tar.addfile(info, f)
elif transferable.is_dir():
tar.add(str(transferable), arcname=f"{destination.rstrip('/')}/{transferable.name}")
else:
raise TypeError(f"Path {transferable} is neither a file nor directory")
else:
raise TypeError("source must be bytes or Path")
return buf.getvalue()
108 changes: 0 additions & 108 deletions core/tests/test_core.py
Loading
Loading