Add type annotations for postgres.

Author: tillahoffmann

postgres/testcontainers/postgres/__init__.py

@@ -11,7 +11,7 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 import os
-
+from typing import Optional
 from testcontainers.core.generic import DbContainer
 
 
@@ -39,13 +39,9 @@ class PostgresContainer(DbContainer):
     POSTGRES_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "test")
     POSTGRES_DB = os.environ.get("POSTGRES_DB", "test")
 
-    def __init__(self,
-                 image="postgres:latest",
-                 port=5432, user=None,
-                 password=None,
-                 dbname=None,
-                 driver="psycopg2",
-                 **kwargs):
+    def __init__(self, image: str = "postgres:latest", port: int = 5432, user: Optional[str] = None,
+                 password: Optional[str] = None, dbname: Optional[str] = None,
+                 driver: str = "psycopg2", **kwargs) -> None:
         super(PostgresContainer, self).__init__(image=image, **kwargs)
         self.POSTGRES_USER = user or self.POSTGRES_USER
         self.POSTGRES_PASSWORD = password or self.POSTGRES_PASSWORD
@@ -55,15 +51,14 @@ def __init__(self,
 
         self.with_exposed_ports(self.port_to_expose)
 
-    def _configure(self):
+    def _configure(self) -> None:
         self.with_env("POSTGRES_USER", self.POSTGRES_USER)
         self.with_env("POSTGRES_PASSWORD", self.POSTGRES_PASSWORD)
         self.with_env("POSTGRES_DB", self.POSTGRES_DB)
 
-    def get_connection_url(self, host=None):
-        return super()._create_connection_url(dialect="postgresql+{}".format(self.driver),
-                                              username=self.POSTGRES_USER,
-                                              password=self.POSTGRES_PASSWORD,
-                                              db_name=self.POSTGRES_DB,
-                                              host=host,
-                                              port=self.port_to_expose)
+    def get_connection_url(self, host=None) -> str:
+        return super()._create_connection_url(
+            dialect="postgresql+{}".format(self.driver), username=self.POSTGRES_USER,
+            password=self.POSTGRES_PASSWORD, db_name=self.POSTGRES_DB, host=host,
+            port=self.port_to_expose,
+        )