Add type annotations for core package.

Author: tillahoffmann

compose/testcontainers/compose/__init__.py

@@ -12,7 +12,7 @@
 from testcontainers.core.exceptions import NoSuchPortExposed
 
 
-class DockerCompose(object):
+class DockerCompose:
     """
     Manage docker compose environments.
 

core/testcontainers/core/container.py

@@ -1,5 +1,6 @@
-import os
 from docker.models.containers import Container
+import os
+from typing import Iterable, Optional, Tuple
 
 from .waiting_utils import wait_container_is_ready
 from .docker_client import DockerClient
@@ -9,7 +10,7 @@
 logger = setup_logger(__name__)
 
 
-class DockerContainer(object):
+class DockerContainer:
     """
     Basic container object to spin up Docker instances.
 
@@ -21,7 +22,7 @@ class DockerContainer(object):
         >>> with DockerContainer("hello-world") as container:
         ...    delay = wait_for_logs(container, "Hello from Docker!")
     """
-    def __init__(self, image, docker_client_kw: dict = None, **kwargs):
+    def __init__(self, image: str, docker_client_kw: Optional[dict] = None, **kwargs) -> None:
         self.env = {}
         self.ports = {}
         self.volumes = {}
@@ -36,13 +37,12 @@ def with_env(self, key: str, value: str) -> 'DockerContainer':
         self.env[key] = value
         return self
 
-    def with_bind_ports(self, container: int,
-                        host: int = None) -> 'DockerContainer':
+    def with_bind_ports(self, container: int, host: int = None) -> 'DockerContainer':
         self.ports[container] = host
         return self
 
-    def with_exposed_ports(self, *ports) -> 'DockerContainer':
-        for port in list(ports):
+    def with_exposed_ports(self, *ports: Iterable[int]) -> 'DockerContainer':
+        for port in ports:
             self.ports[port] = None
         return self
 
@@ -55,25 +55,20 @@ def maybe_emulate_amd64(self) -> 'DockerContainer':
             return self.with_kwargs(platform='linux/amd64')
         return self
 
-    def start(self):
+    def start(self) -> 'DockerContainer':
         logger.info("Pulling image %s", self.image)
         docker_client = self.get_docker_client()
-        self._container = docker_client.run(self.image,
-                                            command=self._command,
-                                            detach=True,
-                                            environment=self.env,
-                                            ports=self.ports,
-                                            name=self._name,
-                                            volumes=self.volumes,
-                                            **self._kwargs
-                                            )
+        self._container = docker_client.run(
+            self.image, command=self._command, detach=True, environment=self.env, ports=self.ports,
+            name=self._name, volumes=self.volumes, **self._kwargs
+        )
         logger.info("Container started: %s", self._container.short_id)
         return self
 
-    def stop(self, force=True, delete_volume=True):
+    def stop(self, force=True, delete_volume=True) -> None:
         self.get_wrapped_container().remove(force=force, v=delete_volume)
 
-    def __enter__(self):
+    def __enter__(self) -> 'DockerContainer':
         return self.start()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -110,7 +105,7 @@ def get_container_host_ip(self) -> str:
         return host
 
     @wait_container_is_ready()
-    def get_exposed_port(self, port) -> str:
+    def get_exposed_port(self, port: int) -> str:
         mapped_port = self.get_docker_client().port(self._container.id, port)
         if inside_container():
             gateway_ip = self.get_docker_client().gateway_ip(self._container.id)
@@ -128,9 +123,7 @@ def with_name(self, name: str) -> 'DockerContainer':
         self._name = name
         return self
 
-    def with_volume_mapping(self, host: str, container: str,
-                            mode: str = 'ro') -> 'DockerContainer':
-        # '/home/user1/': {'bind': '/mnt/vol2', 'mode': 'rw'}
+    def with_volume_mapping(self, host: str, container: str, mode: str = 'ro') -> 'DockerContainer':
         mapping = {'bind': container, 'mode': mode}
         self.volumes[host] = mapping
         return self
@@ -141,12 +134,12 @@ def get_wrapped_container(self) -> Container:
     def get_docker_client(self) -> DockerClient:
         return self._docker
 
-    def get_logs(self):
+    def get_logs(self) -> Tuple[str, str]:
         if not self._container:
-            raise ContainerStartException("Container should be started before")
+            raise ContainerStartException("Container should be started before getting logs")
         return self._container.logs(stderr=False), self._container.logs(stdout=False)
 
-    def exec(self, command):
+    def exec(self, command) -> Tuple[int, str]:
         if not self._container:
-            raise ContainerStartException("Container should be started before")
+            raise ContainerStartException("Container should be started before executing a command")
         return self.get_wrapped_container().exec_run(command)

core/testcontainers/core/docker_client.py

@@ -11,19 +11,21 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 import atexit
-import os
-import urllib
 import docker
 from docker.errors import NotFound
-from docker.models.containers import Container
+from docker.models.containers import Container, ContainerCollection
+import functools as ft
+import os
+from typing import List, Optional, Union
+import urllib
 
 from .utils import default_gateway_ip, inside_container, setup_logger
 
 
 LOGGER = setup_logger(__name__)
 
 
-def _stop_container(container):
+def _stop_container(container: Container) -> None:
     try:
         container.stop()
     except NotFound:
@@ -33,53 +35,63 @@ def _stop_container(container):
                        container.image, ex)
 
 
-class DockerClient(object):
-    def __init__(self, **kwargs):
+class DockerClient:
+    """
+    Thin wrapper around :class:`docker.DockerClient` for a more functional interface.
+    """
+    def __init__(self, **kwargs) -> None:
         self.client = docker.from_env(**kwargs)
 
-    def run(self, image: str,
-            command: str = None,
-            environment: dict = None,
-            ports: dict = None,
-            detach: bool = False,
-            stdout: bool = True,
-            stderr: bool = False,
-            remove: bool = False, **kwargs) -> Container:
-        container = self.client.containers.run(image,
-                                               command=command,
-                                               stdout=stdout,
-                                               stderr=stderr,
-                                               remove=remove,
-                                               detach=detach,
-                                               environment=environment,
-                                               ports=ports,
-                                               **kwargs)
-        atexit.register(_stop_container, container)
-
+    @ft.wraps(ContainerCollection.run)
+    def run(self, image: str, command: Union[str, List[str]] = None,
+            environment: Optional[dict] = None, ports: Optional[dict] = None,
+            detach: bool = False, stdout: bool = True, stderr: bool = False, remove: bool = False,
+            **kwargs) -> Container:
+        container = self.client.containers.run(
+            image, command=command, stdout=stdout, stderr=stderr, remove=remove, detach=detach,
+            environment=environment, ports=ports, **kwargs
+        )
+        if detach:
+            atexit.register(_stop_container, container)
         return container
 
-    def port(self, container_id, port):
+    def port(self, container_id: str, port: int) -> int:
+        """
+        Lookup the public-facing port that is NAT-ed to :code:`port`.
+        """
         port_mappings = self.client.api.port(container_id, port)
         if not port_mappings:
             raise ConnectionError(f'port mapping for container {container_id} and port {port} is '
                                   'not available')
         return port_mappings[0]["HostPort"]
 
-    def get_container(self, container_id):
+    def get_container(self, container_id: str) -> Container:
+        """
+        Get the container with a given identifier.
+        """
         containers = self.client.api.containers(filters={'id': container_id})
         if not containers:
             raise RuntimeError(f'could not get container with id {container_id}')
         return containers[0]
 
-    def bridge_ip(self, container_id):
+    def bridge_ip(self, container_id: str) -> str:
+        """
+        Get the bridge ip address for a container.
+        """
         container = self.get_container(container_id)
         return container['NetworkSettings']['Networks']['bridge']['IPAddress']
 
-    def gateway_ip(self, container_id):
+    def gateway_ip(self, container_id: str) -> str:
+        """
+        Get the gateway ip address for a container.
+        """
         container = self.get_container(container_id)
         return container['NetworkSettings']['Networks']['bridge']['Gateway']
 
-    def host(self):
+    def host(self) -> str:
+        """
+        Get the hostname or ip address of the docker host.
+        """
         # https://github.com/testcontainers/testcontainers-go/blob/dd76d1e39c654433a3d80429690d07abcec04424/docker.go#L644
         # if os env TC_HOST is set, use it
         host = os.environ.get('TC_HOST')

core/testcontainers/core/exceptions.py

@@ -16,5 +16,5 @@ class ContainerStartException(RuntimeError):
     pass
 
 
-class NoSuchPortExposed(Exception):
+class NoSuchPortExposed(RuntimeError):
     pass

core/testcontainers/core/generic.py

@@ -10,8 +10,10 @@
 #    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 #    License for the specific language governing permissions and limitations
 #    under the License.
+from typing import Optional
 
 from .container import DockerContainer
+from .exceptions import ContainerStartException
 from .waiting_utils import wait_container_is_ready
 
 ADDITIONAL_TRANSIENT_ERRORS = []
@@ -23,24 +25,24 @@
 
 
 class DbContainer(DockerContainer):
-    def __init__(self, image, **kwargs):
-        super(DbContainer, self).__init__(image, **kwargs)
-
+    """
+    Generic database container.
+    """
     @wait_container_is_ready(*ADDITIONAL_TRANSIENT_ERRORS)
-    def _connect(self):
+    def _connect(self) -> None:
         import sqlalchemy
         engine = sqlalchemy.create_engine(self.get_connection_url())
         engine.connect()
 
-    def get_connection_url(self):
+    def get_connection_url(self) -> str:
         raise NotImplementedError
 
-    def _create_connection_url(self, dialect, username, password,
-                               host=None, port=None, db_name=None):
+    def _create_connection_url(self, dialect: str, username: str, password: str,
+                               host: Optional[str] = None, port: Optional[int] = None,
+                               db_name: Optional[str] = None) -> str:
         if self._container is None:
-            raise RuntimeError("container has not been started")
-        if not host:
-            host = self.get_container_host_ip()
+            raise ContainerStartException("container has not been started")
+        host = host or self.get_container_host_ip()
         port = self.get_exposed_port(port)
         url = "{dialect}://{username}:{password}@{host}:{port}".format(
             dialect=dialect, username=username, password=password, host=host, port=port
@@ -49,11 +51,11 @@ def _create_connection_url(self, dialect, username, password,
             url += '/' + db_name
         return url
 
-    def start(self):
+    def start(self) -> 'DbContainer':
         self._configure()
         super().start()
         self._connect()
         return self
 
-    def _configure(self):
+    def _configure(self) -> None:
         raise NotImplementedError

core/testcontainers/core/utils.py

@@ -9,7 +9,7 @@
 WIN = "win"
 
 
-def setup_logger(name):
+def setup_logger(name: str) -> logging.Logger:
     logger = logging.getLogger(name)
     logger.setLevel(logging.INFO)
     handler = logging.StreamHandler()
@@ -18,7 +18,7 @@ def setup_logger(name):
     return logger
 
 
-def os_name():
+def os_name() -> str:
     pl = sys.platform
     if pl == "linux" or pl == "linux2":
         return LINUX
@@ -28,23 +28,23 @@ def os_name():
         return WIN
 
 
-def is_mac():
+def is_mac() -> bool:
     return MAC == os_name()
 
 
-def is_linux():
+def is_linux() -> bool:
     return LINUX == os_name()
 
 
-def is_windows():
+def is_windows() -> bool:
     return WIN == os_name()
 
 
-def is_arm():
+def is_arm() -> bool:
     return platform.machine() in ('arm64', 'aarch64')
 
 
-def inside_container():
+def inside_container() -> bool:
     """
     Returns true if we are running inside a container.
 
@@ -53,7 +53,7 @@ def inside_container():
     return os.path.exists('/.dockerenv')
 
 
-def default_gateway_ip():
+def default_gateway_ip() -> str:
     """
     Returns gateway IP address of the host that testcontainer process is
     running on