Skip to content

Commit fce23a4

Browse files
tillahoffmanntillahoffmann
committed
Add type annotations for core package.
1 parent 383b12e commit fce23a4

7 files changed

Lines changed: 99 additions & 88 deletions

File tree

‎compose/testcontainers/compose/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from testcontainers.core.exceptions import NoSuchPortExposed
1313

1414

15-
class DockerCompose(object):
15+
class DockerCompose:
1616
"""
1717
Manage docker compose environments.
1818

‎core/testcontainers/core/container.py‎

Lines changed: 20 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
1-
import os
21
from docker.models.containers import Container
2+
import os
3+
from typing import Iterable, Optional, Tuple
34

45
from .waiting_utils import wait_container_is_ready
56
from .docker_client import DockerClient
@@ -9,7 +10,7 @@
910
logger = setup_logger(__name__)
1011

1112

12-
class DockerContainer(object):
13+
class DockerContainer:
1314
"""
1415
Basic container object to spin up Docker instances.
1516
@@ -21,7 +22,7 @@ class DockerContainer(object):
2122
>>> with DockerContainer("hello-world") as container:
2223
... delay = wait_for_logs(container, "Hello from Docker!")
2324
"""
24-
def __init__(self, image, docker_client_kw: dict = None, **kwargs):
25+
def __init__(self, image: str, docker_client_kw: Optional[dict] = None, **kwargs) -> None:
2526
self.env = {}
2627
self.ports = {}
2728
self.volumes = {}
@@ -36,13 +37,12 @@ def with_env(self, key: str, value: str) -> 'DockerContainer':
3637
self.env[key] = value
3738
return self
3839

39-
def with_bind_ports(self, container: int,
40-
host: int = None) -> 'DockerContainer':
40+
def with_bind_ports(self, container: int, host: int = None) -> 'DockerContainer':
4141
self.ports[container] = host
4242
return self
4343

44-
def with_exposed_ports(self, *ports) -> 'DockerContainer':
45-
for port in list(ports):
44+
def with_exposed_ports(self, *ports: Iterable[int]) -> 'DockerContainer':
45+
for port in ports:
4646
self.ports[port] = None
4747
return self
4848

@@ -55,25 +55,20 @@ def maybe_emulate_amd64(self) -> 'DockerContainer':
5555
return self.with_kwargs(platform='linux/amd64')
5656
return self
5757

58-
def start(self):
58+
def start(self) -> 'DockerContainer':
5959
logger.info("Pulling image %s", self.image)
6060
docker_client = self.get_docker_client()
61-
self._container = docker_client.run(self.image,
62-
command=self._command,
63-
detach=True,
64-
environment=self.env,
65-
ports=self.ports,
66-
name=self._name,
67-
volumes=self.volumes,
68-
**self._kwargs
69-
)
61+
self._container = docker_client.run(
62+
self.image, command=self._command, detach=True, environment=self.env, ports=self.ports,
63+
name=self._name, volumes=self.volumes, **self._kwargs
64+
)
7065
logger.info("Container started: %s", self._container.short_id)
7166
return self
7267

73-
def stop(self, force=True, delete_volume=True):
68+
def stop(self, force=True, delete_volume=True) -> None:
7469
self.get_wrapped_container().remove(force=force, v=delete_volume)
7570

76-
def __enter__(self):
71+
def __enter__(self) -> 'DockerContainer':
7772
return self.start()
7873

7974
def __exit__(self, exc_type, exc_val, exc_tb):
@@ -110,7 +105,7 @@ def get_container_host_ip(self) -> str:
110105
return host
111106

112107
@wait_container_is_ready()
113-
def get_exposed_port(self, port) -> str:
108+
def get_exposed_port(self, port: int) -> str:
114109
mapped_port = self.get_docker_client().port(self._container.id, port)
115110
if inside_container():
116111
gateway_ip = self.get_docker_client().gateway_ip(self._container.id)
@@ -128,9 +123,7 @@ def with_name(self, name: str) -> 'DockerContainer':
128123
self._name = name
129124
return self
130125

131-
def with_volume_mapping(self, host: str, container: str,
132-
mode: str = 'ro') -> 'DockerContainer':
133-
# '/home/user1/': {'bind': '/mnt/vol2', 'mode': 'rw'}
126+
def with_volume_mapping(self, host: str, container: str, mode: str = 'ro') -> 'DockerContainer':
134127
mapping = {'bind': container, 'mode': mode}
135128
self.volumes[host] = mapping
136129
return self
@@ -141,12 +134,12 @@ def get_wrapped_container(self) -> Container:
141134
def get_docker_client(self) -> DockerClient:
142135
return self._docker
143136

144-
def get_logs(self):
137+
def get_logs(self) -> Tuple[str, str]:
145138
if not self._container:
146-
raise ContainerStartException("Container should be started before")
139+
raise ContainerStartException("Container should be started before getting logs")
147140
return self._container.logs(stderr=False), self._container.logs(stdout=False)
148141

149-
def exec(self, command):
142+
def exec(self, command) -> Tuple[int, str]:
150143
if not self._container:
151-
raise ContainerStartException("Container should be started before")
144+
raise ContainerStartException("Container should be started before executing a command")
152145
return self.get_wrapped_container().exec_run(command)

‎core/testcontainers/core/docker_client.py‎

Lines changed: 42 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -11,19 +11,21 @@
1111
# License for the specific language governing permissions and limitations
1212
# under the License.
1313
import atexit
14-
import os
15-
import urllib
1614
import docker
1715
from docker.errors import NotFound
18-
from docker.models.containers import Container
16+
from docker.models.containers import Container, ContainerCollection
17+
import functools as ft
18+
import os
19+
from typing import List, Optional, Union
20+
import urllib
1921

2022
from .utils import default_gateway_ip, inside_container, setup_logger
2123

2224

2325
LOGGER = setup_logger(__name__)
2426

2527

26-
def _stop_container(container):
28+
def _stop_container(container: Container) -> None:
2729
try:
2830
container.stop()
2931
except NotFound:
@@ -33,53 +35,63 @@ def _stop_container(container):
3335
container.image, ex)
3436

3537

36-
class DockerClient(object):
37-
def __init__(self, **kwargs):
38+
class DockerClient:
39+
"""
40+
Thin wrapper around :class:`docker.DockerClient` for a more functional interface.
41+
"""
42+
def __init__(self, **kwargs) -> None:
3843
self.client = docker.from_env(**kwargs)
3944

40-
def run(self, image: str,
41-
command: str = None,
42-
environment: dict = None,
43-
ports: dict = None,
44-
detach: bool = False,
45-
stdout: bool = True,
46-
stderr: bool = False,
47-
remove: bool = False, **kwargs) -> Container:
48-
container = self.client.containers.run(image,
49-
command=command,
50-
stdout=stdout,
51-
stderr=stderr,
52-
remove=remove,
53-
detach=detach,
54-
environment=environment,
55-
ports=ports,
56-
**kwargs)
57-
atexit.register(_stop_container, container)
58-
45+
@ft.wraps(ContainerCollection.run)
46+
def run(self, image: str, command: Union[str, List[str]] = None,
47+
environment: Optional[dict] = None, ports: Optional[dict] = None,
48+
detach: bool = False, stdout: bool = True, stderr: bool = False, remove: bool = False,
49+
**kwargs) -> Container:
50+
container = self.client.containers.run(
51+
image, command=command, stdout=stdout, stderr=stderr, remove=remove, detach=detach,
52+
environment=environment, ports=ports, **kwargs
53+
)
54+
if detach:
55+
atexit.register(_stop_container, container)
5956
return container
6057

61-
def port(self, container_id, port):
58+
def port(self, container_id: str, port: int) -> int:
59+
"""
60+
Lookup the public-facing port that is NAT-ed to :code:`port`.
61+
"""
6262
port_mappings = self.client.api.port(container_id, port)
6363
if not port_mappings:
6464
raise ConnectionError(f'port mapping for container {container_id} and port {port} is '
6565
'not available')
6666
return port_mappings[0]["HostPort"]
6767

68-
def get_container(self, container_id):
68+
def get_container(self, container_id: str) -> Container:
69+
"""
70+
Get the container with a given identifier.
71+
"""
6972
containers = self.client.api.containers(filters={'id': container_id})
7073
if not containers:
7174
raise RuntimeError(f'could not get container with id {container_id}')
7275
return containers[0]
7376

74-
def bridge_ip(self, container_id):
77+
def bridge_ip(self, container_id: str) -> str:
78+
"""
79+
Get the bridge ip address for a container.
80+
"""
7581
container = self.get_container(container_id)
7682
return container['NetworkSettings']['Networks']['bridge']['IPAddress']
7783

78-
def gateway_ip(self, container_id):
84+
def gateway_ip(self, container_id: str) -> str:
85+
"""
86+
Get the gateway ip address for a container.
87+
"""
7988
container = self.get_container(container_id)
8089
return container['NetworkSettings']['Networks']['bridge']['Gateway']
8190

82-
def host(self):
91+
def host(self) -> str:
92+
"""
93+
Get the hostname or ip address of the docker host.
94+
"""
8395
# https://github.com/testcontainers/testcontainers-go/blob/dd76d1e39c654433a3d80429690d07abcec04424/docker.go#L644
8496
# if os env TC_HOST is set, use it
8597
host = os.environ.get('TC_HOST')

‎core/testcontainers/core/exceptions.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,5 +16,5 @@ class ContainerStartException(RuntimeError):
1616
pass
1717

1818

19-
class NoSuchPortExposed(Exception):
19+
class NoSuchPortExposed(RuntimeError):
2020
pass

‎core/testcontainers/core/generic.py‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,10 @@
1010
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
1111
# License for the specific language governing permissions and limitations
1212
# under the License.
13+
from typing import Optional
1314

1415
from .container import DockerContainer
16+
from .exceptions import ContainerStartException
1517
from .waiting_utils import wait_container_is_ready
1618

1719
ADDITIONAL_TRANSIENT_ERRORS = []
@@ -23,24 +25,24 @@
2325

2426

2527
class DbContainer(DockerContainer):
26-
def __init__(self, image, **kwargs):
27-
super(DbContainer, self).__init__(image, **kwargs)
28-
28+
"""
29+
Generic database container.
30+
"""
2931
@wait_container_is_ready(*ADDITIONAL_TRANSIENT_ERRORS)
30-
def _connect(self):
32+
def _connect(self) -> None:
3133
import sqlalchemy
3234
engine = sqlalchemy.create_engine(self.get_connection_url())
3335
engine.connect()
3436

35-
def get_connection_url(self):
37+
def get_connection_url(self) -> str:
3638
raise NotImplementedError
3739

38-
def _create_connection_url(self, dialect, username, password,
39-
host=None, port=None, db_name=None):
40+
def _create_connection_url(self, dialect: str, username: str, password: str,
41+
host: Optional[str] = None, port: Optional[int] = None,
42+
db_name: Optional[str] = None) -> str:
4043
if self._container is None:
41-
raise RuntimeError("container has not been started")
42-
if not host:
43-
host = self.get_container_host_ip()
44+
raise ContainerStartException("container has not been started")
45+
host = host or self.get_container_host_ip()
4446
port = self.get_exposed_port(port)
4547
url = "{dialect}://{username}:{password}@{host}:{port}".format(
4648
dialect=dialect, username=username, password=password, host=host, port=port
@@ -49,11 +51,11 @@ def _create_connection_url(self, dialect, username, password,
4951
url += '/' + db_name
5052
return url
5153

52-
def start(self):
54+
def start(self) -> 'DbContainer':
5355
self._configure()
5456
super().start()
5557
self._connect()
5658
return self
5759

58-
def _configure(self):
60+
def _configure(self) -> None:
5961
raise NotImplementedError

‎core/testcontainers/core/utils.py‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
WIN = "win"
1010

1111

12-
def setup_logger(name):
12+
def setup_logger(name: str) -> logging.Logger:
1313
logger = logging.getLogger(name)
1414
logger.setLevel(logging.INFO)
1515
handler = logging.StreamHandler()
@@ -18,7 +18,7 @@ def setup_logger(name):
1818
return logger
1919

2020

21-
def os_name():
21+
def os_name() -> str:
2222
pl = sys.platform
2323
if pl == "linux" or pl == "linux2":
2424
return LINUX
@@ -28,23 +28,23 @@ def os_name():
2828
return WIN
2929

3030

31-
def is_mac():
31+
def is_mac() -> bool:
3232
return MAC == os_name()
3333

3434

35-
def is_linux():
35+
def is_linux() -> bool:
3636
return LINUX == os_name()
3737

3838

39-
def is_windows():
39+
def is_windows() -> bool:
4040
return WIN == os_name()
4141

4242

43-
def is_arm():
43+
def is_arm() -> bool:
4444
return platform.machine() in ('arm64', 'aarch64')
4545

4646

47-
def inside_container():
47+
def inside_container() -> bool:
4848
"""
4949
Returns true if we are running inside a container.
5050
@@ -53,7 +53,7 @@ def inside_container():
5353
return os.path.exists('/.dockerenv')
5454

5555

56-
def default_gateway_ip():
56+
def default_gateway_ip() -> str:
5757
"""
5858
Returns gateway IP address of the host that testcontainer process is
5959
running on

0 commit comments

Comments
 (0)