Add type annotations for mongodb.

Author: tillahoffmann

mongodb/testcontainers/mongodb/__init__.py

@@ -11,7 +11,7 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 import os
-
+from pymongo import MongoClient
 from testcontainers.core.generic import DbContainer
 from testcontainers.core.waiting_utils import wait_container_is_ready
 
@@ -50,21 +50,18 @@ class MongoDbContainer(DbContainer):
     MONGO_INITDB_ROOT_PASSWORD = os.environ.get("MONGO_INITDB_ROOT_PASSWORD", "test")
     MONGO_DB = os.environ.get("MONGO_DB", "test")
 
-    def __init__(self,
-                 image: str = "mongo:latest",
-                 port_to_expose: int = 27017,
-                 **kwargs):
+    def __init__(self, image: str = "mongo:latest", port_to_expose: int = 27017, **kwargs) -> None:
         super(MongoDbContainer, self).__init__(image=image, **kwargs)
         self.command = "mongo"
         self.port_to_expose = port_to_expose
         self.with_exposed_ports(self.port_to_expose)
 
-    def _configure(self):
+    def _configure(self) -> None:
         self.with_env("MONGO_INITDB_ROOT_USERNAME", self.MONGO_INITDB_ROOT_USERNAME)
         self.with_env("MONGO_INITDB_ROOT_PASSWORD", self.MONGO_INITDB_ROOT_PASSWORD)
         self.with_env("MONGO_DB", self.MONGO_DB)
 
-    def get_connection_url(self):
+    def get_connection_url(self) -> str:
         return self._create_connection_url(
             dialect='mongodb',
             username=self.MONGO_INITDB_ROOT_USERNAME,
@@ -73,9 +70,8 @@ def get_connection_url(self):
         )
 
     @wait_container_is_ready()
-    def _connect(self):
-        from pymongo import MongoClient
+    def _connect(self) -> MongoClient:
         return MongoClient(self.get_connection_url())
 
-    def get_connection_client(self):
+    def get_connection_client(self) -> MongoClient:
         return self._connect()