Skip to content

Commit d5b8553

Browse files
authoredApr 3, 2024··
fix(core): Improve typing for common container usage scenarios (#523)
Improves type hints for type checking in a common use cases: ```python with MySqlContainer("mysql:8").with_env("some", "value") as mysql: url = mysql.get_connection_url() # get_connection_url would previously be an unknown member here ``` And, also improved type hinting for the custom `DockerClient`'s `run` command, where the linter no longer reports an error due to missing parameter types: ```python DockerClient.run("nginx") # Previously this would report "Argument missing for parameter "image" ```
1 parent 0fb4aef commit d5b8553

File tree

2 files changed

+28
-13
lines changed

2 files changed

+28
-13
lines changed
 

‎core/testcontainers/core/container.py

+13-11
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from typing import TYPE_CHECKING, Optional
55

66
import docker.errors
7+
from typing_extensions import Self
78

89
from testcontainers.core.config import (
910
RYUK_DISABLED,
@@ -53,29 +54,29 @@ def __init__(
5354
self._name = None
5455
self._kwargs = kwargs
5556

56-
def with_env(self, key: str, value: str) -> "DockerContainer":
57+
def with_env(self, key: str, value: str) -> Self:
5758
self.env[key] = value
5859
return self
5960

60-
def with_bind_ports(self, container: int, host: Optional[int] = None) -> "DockerContainer":
61+
def with_bind_ports(self, container: int, host: Optional[int] = None) -> Self:
6162
self.ports[container] = host
6263
return self
6364

64-
def with_exposed_ports(self, *ports: int) -> "DockerContainer":
65+
def with_exposed_ports(self, *ports: int) -> Self:
6566
for port in ports:
6667
self.ports[port] = None
6768
return self
6869

69-
def with_kwargs(self, **kwargs) -> "DockerContainer":
70+
def with_kwargs(self, **kwargs) -> Self:
7071
self._kwargs = kwargs
7172
return self
7273

73-
def maybe_emulate_amd64(self) -> "DockerContainer":
74+
def maybe_emulate_amd64(self) -> Self:
7475
if is_arm():
7576
return self.with_kwargs(platform="linux/amd64")
7677
return self
7778

78-
def start(self):
79+
def start(self) -> Self:
7980
if not RYUK_DISABLED and self.image != RYUK_IMAGE:
8081
logger.debug("Creating Ryuk container")
8182
Reaper.get_instance()
@@ -95,10 +96,11 @@ def start(self):
9596
return self
9697

9798
def stop(self, force=True, delete_volume=True) -> None:
98-
self._container.remove(force=force, v=delete_volume)
99+
if self._container:
100+
self._container.remove(force=force, v=delete_volume)
99101
self.get_docker_client().client.close()
100102

101-
def __enter__(self):
103+
def __enter__(self) -> Self:
102104
return self.start()
103105

104106
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
@@ -138,15 +140,15 @@ def get_exposed_port(self, port: int) -> str:
138140
return port
139141
return mapped_port
140142

141-
def with_command(self, command: str) -> "DockerContainer":
143+
def with_command(self, command: str) -> Self:
142144
self._command = command
143145
return self
144146

145-
def with_name(self, name: str) -> "DockerContainer":
147+
def with_name(self, name: str) -> Self:
146148
self._name = name
147149
return self
148150

149-
def with_volume_mapping(self, host: str, container: str, mode: str = "ro") -> "DockerContainer":
151+
def with_volume_mapping(self, host: str, container: str, mode: str = "ro") -> Self:
150152
mapping = {"bind": container, "mode": mode}
151153
self.volumes[host] = mapping
152154
return self

‎core/testcontainers/core/docker_client.py

+15-2
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,11 @@
1818
import urllib.parse
1919
from os.path import exists
2020
from pathlib import Path
21-
from typing import Optional, Union
21+
from typing import Callable, Optional, TypeVar, Union
2222

2323
import docker
2424
from docker.models.containers import Container, ContainerCollection
25+
from typing_extensions import ParamSpec
2526

2627
from testcontainers.core.labels import SESSION_ID, create_labels
2728
from testcontainers.core.utils import default_gateway_ip, inside_container, setup_logger
@@ -30,6 +31,18 @@
3031
TC_FILE = ".testcontainers.properties"
3132
TC_GLOBAL = Path.home() / TC_FILE
3233

34+
_P = ParamSpec("_P")
35+
_T = TypeVar("_T")
36+
37+
38+
def _wrapped_container_collection(function: Callable[_P, _T]) -> Callable[_P, _T]:
39+
40+
@ft.wraps(ContainerCollection.run)
41+
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
42+
return function(*args, **kwargs)
43+
44+
return wrapper
45+
3346

3447
class DockerClient:
3548
"""
@@ -48,7 +61,7 @@ def __init__(self, **kwargs) -> None:
4861
self.client.api.headers["x-tc-sid"] = SESSION_ID
4962
self.client.api.headers["User-Agent"] = "tc-python/" + importlib.metadata.version("testcontainers")
5063

51-
@ft.wraps(ContainerCollection.run)
64+
@_wrapped_container_collection
5265
def run(
5366
self,
5467
image: str,

0 commit comments

Comments
 (0)
Please sign in to comment.