refactor(typing): annotate self-returning methods with Self (#1008)

Co-authored-by: Roy Moore <roy@moore.co.il>

Author: ThanasisMattas

src/testcontainers/community/azurite/__init__.py

@@ -14,6 +14,8 @@
 import os
 from typing import Optional
 
+from typing_extensions import Self
+
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.utils import raise_for_deprecated_parameter
 from testcontainers.core.wait_strategies import PortWaitStrategy
@@ -217,7 +219,7 @@ def __get_external_connection_string(self) -> str:
 
         return connection_string
 
-    def start(self) -> "AzuriteContainer":
+    def start(self) -> Self:
         super().start()
         self._connect()
         return self

src/testcontainers/community/chroma/__init__.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING
 
 from requests import ConnectionError, get
+from typing_extensions import Self
 
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.utils import raise_for_deprecated_parameter
@@ -73,7 +74,7 @@ def _healthcheck(self) -> None:
         response: Response = get(url)
         response.raise_for_status()
 
-    def start(self) -> "ChromaContainer":
+    def start(self) -> Self:
         """This method starts the Chroma container and runs the healthcheck
         to verify that the container is ready to use."""
         super().start()

src/testcontainers/community/generic/server.py

@@ -1,6 +1,7 @@
 from typing import Union
 
 import httpx
+from typing_extensions import Self
 
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.exceptions import ContainerStartException
@@ -53,7 +54,7 @@ def _create_connection_url(self) -> str:
         url = f"http://{host}:{exposed_port}"
         return url
 
-    def start(self) -> "ServerContainer":
+    def start(self) -> Self:
         super().start()
         self._connect()
         return self

src/testcontainers/community/generic/sql.py

@@ -2,6 +2,8 @@
 from typing import Any, Optional
 from urllib.parse import quote, urlencode
 
+from typing_extensions import Self
+
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.exceptions import ContainerStartException
 from testcontainers.core.waiting_utils import WaitStrategy
@@ -83,7 +85,7 @@ def _create_connection_url(
 
         return url
 
-    def start(self) -> "SqlContainer":
+    def start(self) -> Self:
         """
         Start the database container and perform initialization.
 

src/testcontainers/community/influxdb/base.py

@@ -15,6 +15,7 @@
 
 from requests import get
 from requests.exceptions import ConnectionError, ReadTimeout
+from typing_extensions import Self
 
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_container_is_ready
@@ -74,7 +75,7 @@ def get_influxdb_version(self) -> str:
 
         return self._health_check().get("version")
 
-    def start(self) -> "InfluxDbContainer":
+    def start(self) -> Self:
         """
         Spawns a container of the InfluxDB Docker image, ready to be used.
         """

src/testcontainers/community/influxdb/version1.py

@@ -14,6 +14,7 @@
 from typing import Optional
 
 from influxdb import InfluxDBClient
+from typing_extensions import Self
 
 from .base import InfluxDbContainer
 
@@ -58,7 +59,7 @@ def get_client(self, **client_kwargs):
 
         return InfluxDBClient(self.get_container_host_ip(), self.get_exposed_port(self.container_port), **client_kwargs)
 
-    def start(self) -> "InfluxDb1Container":
+    def start(self) -> Self:
         """
         Overridden for better typing reason
         """

src/testcontainers/community/influxdb/version2.py

@@ -15,6 +15,7 @@
 from typing import Optional
 
 from influxdb_client import InfluxDBClient, Organization
+from typing_extensions import Self
 
 from .base import InfluxDbContainer
 
@@ -70,7 +71,7 @@ def __init__(
             if env_value:
                 self.with_env(env_key, env_value)
 
-    def start(self) -> "InfluxDb2Container":
+    def start(self) -> Self:
         """
         Overridden for better typing reason
         """

src/testcontainers/community/k3s/__init__.py

@@ -13,6 +13,8 @@
 
 import logging
 
+from typing_extensions import Self
+
 from testcontainers.core.config import testcontainers_config
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
@@ -55,7 +57,7 @@ def __init__(self, image="rancher/k3s:latest", enable_cgroup_mount=True, **kwarg
     def _connect(self) -> None:
         wait_for_logs(self, predicate="Node controller sync successful", timeout=testcontainers_config.timeout)
 
-    def start(self) -> "K3SContainer":
+    def start(self) -> Self:
         super().start()
         self._connect()
         return self

src/testcontainers/community/kafka/__init__.py

@@ -183,7 +183,7 @@ def tc_start(self) -> None:
         )
         self.create_file(data, KafkaContainer.TC_START_SCRIPT)
 
-    def start(self, timeout: int = 30) -> "KafkaContainer":
+    def start(self, timeout: int = 30) -> Self:
         script = KafkaContainer.TC_START_SCRIPT
         command = f'sh -c "while [ ! -f {script} ]; do sleep 0.1; done; sh {script}"'
         self.configure()

src/testcontainers/community/kafka/_redpanda.py

@@ -5,6 +5,8 @@
 from io import BytesIO
 from textwrap import dedent
 
+from typing_extensions import Self
+
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.wait_strategies import LogMessageWaitStrategy
 
@@ -66,7 +68,7 @@ def tc_start(self) -> None:
 
         self.create_file(data, RedpandaContainer.TC_START_SCRIPT)
 
-    def start(self, timeout=10) -> "RedpandaContainer":
+    def start(self, timeout=10) -> Self:
         script = RedpandaContainer.TC_START_SCRIPT
         command = f'-c "while [ ! -f {script} ]; do sleep 0.1; done; sh {script}"'
         self.with_command(command)