fix(core): mypy (#810)

makes mypy happier

---------

Co-authored-by: Roy Moore <[email protected]>

Author: alexanderankin

core/testcontainers/compose/__init__.py

@@ -1,8 +1,15 @@
 # flake8: noqa: F401
 from testcontainers.compose.compose import (
     ComposeContainer,
-    ContainerIsNotRunning,
     DockerCompose,
-    NoSuchPortExposed,
-    PublishedPort,
+    PublishedPortModel,
 )
+from testcontainers.core.exceptions import ContainerIsNotRunning, NoSuchPortExposed
+
+__all__ = [
+    "ComposeContainer",
+    "ContainerIsNotRunning",
+    "DockerCompose",
+    "NoSuchPortExposed",
+    "PublishedPortModel",
+]

core/testcontainers/compose/compose.py

@@ -1,4 +1,4 @@
-from dataclasses import asdict, dataclass, field, fields
+from dataclasses import asdict, dataclass, field, fields, is_dataclass
 from functools import cached_property
 from json import loads
 from logging import warning
@@ -7,6 +7,7 @@
 from re import split
 from subprocess import CompletedProcess
 from subprocess import run as subprocess_run
+from types import TracebackType
 from typing import Any, Callable, Literal, Optional, TypeVar, Union, cast
 from urllib.error import HTTPError, URLError
 from urllib.request import urlopen
@@ -18,35 +19,37 @@
 _WARNINGS = {"DOCKER_COMPOSE_GET_CONFIG": "get_config is experimental, see testcontainers/testcontainers-python#669"}
 
 
-def _ignore_properties(cls: type[_IPT], dict_: any) -> _IPT:
+def _ignore_properties(cls: type[_IPT], dict_: Any) -> _IPT:
     """omits extra fields like @JsonIgnoreProperties(ignoreUnknown = true)
 
     https://gist.github.com/alexanderankin/2a4549ac03554a31bef6eaaf2eaf7fd5"""
     if isinstance(dict_, cls):
         return dict_
+    if not is_dataclass(cls):
+        raise TypeError(f"Expected a dataclass type, got {cls}")
     class_fields = {f.name for f in fields(cls)}
     filtered = {k: v for k, v in dict_.items() if k in class_fields}
-    return cls(**filtered)
+    return cast("_IPT", cls(**filtered))
 
 
 @dataclass
-class PublishedPort:
+class PublishedPortModel:
     """
     Class that represents the response we get from compose when inquiring status
     via `DockerCompose.get_running_containers()`.
     """
 
     URL: Optional[str] = None
-    TargetPort: Optional[str] = None
-    PublishedPort: Optional[str] = None
+    TargetPort: Optional[int] = None
+    PublishedPort: Optional[int] = None
     Protocol: Optional[str] = None
 
-    def normalize(self):
+    def normalize(self) -> "PublishedPortModel":
         url_not_usable = system() == "Windows" and self.URL == "0.0.0.0"
         if url_not_usable:
             self_dict = asdict(self)
             self_dict.update({"URL": "127.0.0.1"})
-            return PublishedPort(**self_dict)
+            return PublishedPortModel(**self_dict)
         return self
 
 
@@ -75,19 +78,19 @@ class ComposeContainer:
     Service: Optional[str] = None
     State: Optional[str] = None
     Health: Optional[str] = None
-    ExitCode: Optional[str] = None
-    Publishers: list[PublishedPort] = field(default_factory=list)
+    ExitCode: Optional[int] = None
+    Publishers: list[PublishedPortModel] = field(default_factory=list)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.Publishers:
-            self.Publishers = [_ignore_properties(PublishedPort, p) for p in self.Publishers]
+            self.Publishers = [_ignore_properties(PublishedPortModel, p) for p in self.Publishers]
 
     def get_publisher(
         self,
         by_port: Optional[int] = None,
         by_host: Optional[str] = None,
-        prefer_ip_version: Literal["IPV4", "IPv6"] = "IPv4",
-    ) -> PublishedPort:
+        prefer_ip_version: Literal["IPv4", "IPv6"] = "IPv4",
+    ) -> PublishedPortModel:
         remaining_publishers = self.Publishers
 
         remaining_publishers = [r for r in remaining_publishers if self._matches_protocol(prefer_ip_version, r)]
@@ -109,8 +112,9 @@ def get_publisher(
         )
 
     @staticmethod
-    def _matches_protocol(prefer_ip_version, r):
-        return (":" in r.URL) is (prefer_ip_version == "IPv6")
+    def _matches_protocol(prefer_ip_version: str, r: PublishedPortModel) -> bool:
+        r_url = r.URL
+        return (r_url is not None and ":" in r_url) is (prefer_ip_version == "IPv6")
 
 
 @dataclass
@@ -164,7 +168,7 @@ class DockerCompose:
                 image: "hello-world"
     """
 
-    context: Union[str, PathLike]
+    context: Union[str, PathLike[str]]
     compose_file_name: Optional[Union[str, list[str]]] = None
     pull: bool = False
     build: bool = False
@@ -175,15 +179,17 @@ class DockerCompose:
     docker_command_path: Optional[str] = None
     profiles: Optional[list[str]] = None
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if isinstance(self.compose_file_name, str):
             self.compose_file_name = [self.compose_file_name]
 
     def __enter__(self) -> "DockerCompose":
         self.start()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+    def __exit__(
+        self, exc_type: Optional[type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
         self.stop(not self.keep_volumes)
 
     def docker_compose_command(self) -> list[str]:
@@ -235,7 +241,7 @@ def start(self) -> None:
 
         self._run_command(cmd=up_cmd)
 
-    def stop(self, down=True) -> None:
+    def stop(self, down: bool = True) -> None:
         """
         Stops the docker compose environment.
         """
@@ -295,7 +301,7 @@ def get_config(
         cmd_output = self._run_command(cmd=config_cmd).stdout
         return cast(dict[str, Any], loads(cmd_output))  # noqa: TC006
 
-    def get_containers(self, include_all=False) -> list[ComposeContainer]:
+    def get_containers(self, include_all: bool = False) -> list[ComposeContainer]:
         """
         Fetch information about running containers via `docker compose ps --format json`.
         Available only in V2 of compose.
@@ -370,17 +376,18 @@ def exec_in_container(
         """
         if not service_name:
             service_name = self.get_container().Service
-        exec_cmd = [*self.compose_command_property, "exec", "-T", service_name, *command]
+        assert service_name
+        exec_cmd: list[str] = [*self.compose_command_property, "exec", "-T", service_name, *command]
         result = self._run_command(cmd=exec_cmd)
 
-        return (result.stdout.decode("utf-8"), result.stderr.decode("utf-8"), result.returncode)
+        return result.stdout.decode("utf-8"), result.stderr.decode("utf-8"), result.returncode
 
     def _run_command(
         self,
         cmd: Union[str, list[str]],
         context: Optional[str] = None,
     ) -> CompletedProcess[bytes]:
-        context = context or self.context
+        context = context or str(self.context)
         return subprocess_run(
             cmd,
             capture_output=True,
@@ -392,7 +399,7 @@ def get_service_port(
         self,
         service_name: Optional[str] = None,
         port: Optional[int] = None,
-    ):
+    ) -> Optional[int]:
         """
         Returns the mapped port for one of the services.
 
@@ -408,13 +415,14 @@ def get_service_port(
         str:
             The mapped port on the host
         """
-        return self.get_container(service_name).get_publisher(by_port=port).normalize().PublishedPort
+        normalize: PublishedPortModel = self.get_container(service_name).get_publisher(by_port=port).normalize()
+        return normalize.PublishedPort
 
     def get_service_host(
         self,
         service_name: Optional[str] = None,
         port: Optional[int] = None,
-    ):
+    ) -> Optional[str]:
         """
         Returns the host for one of the services.
 
@@ -430,13 +438,17 @@ def get_service_host(
         str:
             The hostname for the service
         """
-        return self.get_container(service_name).get_publisher(by_port=port).normalize().URL
+        container: ComposeContainer = self.get_container(service_name)
+        publisher: PublishedPortModel = container.get_publisher(by_port=port)
+        normalize: PublishedPortModel = publisher.normalize()
+        url: Optional[str] = normalize.URL
+        return url
 
     def get_service_host_and_port(
         self,
         service_name: Optional[str] = None,
         port: Optional[int] = None,
-    ):
+    ) -> tuple[Optional[str], Optional[int]]:
         publisher = self.get_container(service_name).get_publisher(by_port=port).normalize()
         return publisher.URL, publisher.PublishedPort
 

core/testcontainers/core/config.py

@@ -7,7 +7,7 @@
 from os import environ
 from os.path import exists
 from pathlib import Path
-from typing import Final, Optional, Union
+from typing import Final, Optional, Union, cast
 
 import docker
 
@@ -39,6 +39,7 @@ def get_docker_socket() -> str:
     try:
         client = docker.from_env()
         socket_path = client.api.get_adapter(client.api.base_url).socket_path
+        socket_path = cast("str", socket_path)
         # return the normalized path as string
         return str(Path(socket_path).absolute())
     except Exception:
@@ -148,6 +149,7 @@ def ryuk_docker_socket(self, value: str) -> None:
 
 __all__ = [
     # Public API of this module:
+    "ConnectionMode",
     "testcontainers_config",
 ]