Typing, docstrings, formatting for mysql_wrapper.py

core/mysql_wrapper.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 import sys
+from typing import Union
 
 sys.path.append("..")
 import mysql.connector
@@ -11,7 +12,18 @@ tasks_db_parameters = {"database": "tasks", **current_db_parameters}
 
 
 class DatabaseWrapper:
-    def __init__(self, connection_parameters):
+    """
+    Wrapper to mantain a connection to the database and launch queries to it
+    in a more convenient way.
+    """
+
+    def __init__(self, connection_parameters: dict) -> None:
+        """
+        Initializes object with parameters and sends a ping to the database
+        to confirm it is working.
+        :param connection_parameters: a dictionary containing the connection
+        parameters for the database.
+        """
         self.host = connection_parameters["host"]
         self.database = connection_parameters["database"]
         self.user = connection_parameters["user"]
@@ -20,7 +32,11 @@ class DatabaseWrapper:
 
         self.ping()
 
-    def connect(self):
+    def connect(self) -> None:
+        """
+        Starts a connection to the database.
+        :return: None
+        """
         self.connection = mysql.connector.connect(
             host=self.host,
             database=self.database,
@@ -29,17 +45,40 @@ class DatabaseWrapper:
             autocommit=False,
         )
 
-    def disconnect(self):
+    def disconnect(self) -> None:
+        """
+        Ends the current connection to the database.
+        :return: None
+        """
         if self.connection.is_connected():
             self.connection.disconnect()
 
-    def ping(self):
+    def ping(self) -> None:
+        """
+        Connects and disconnects to the database to test whether it is up.
+        :return: None
+        """
         self.connect()
         self.disconnect()
 
-    def query(self, query_statement, query_parameters=None, dictionary=False):
+    def query(
+        self,
+        query_statement: str,
+        query_parameters: Union[dict, None] = None,
+        dictionary: bool = False,
+    ) -> mysql.connector.connection.MySQLCursor:
+        """
+        Connects to the database, sends a query and returns the cursor.
+        :param query_statement: string with the query to execute
+        :param query_parameters: parameters to substitute in the query string
+        :param dictionary: whether the parameters are structured as a dict or
+        not
+        :return: cursor for the query
+        """
         self.connect()
-        if self.connection.is_connected():
+        if not self.connection.is_connected():
+            raise Exception("Could not connect to the database.")
+
         try:
             execution_cursor = self.connection.cursor(
                 dictionary=dictionary, buffered=True
@@ -61,12 +100,6 @@ class DatabaseWrapper:
                 ),
             )
 
-        else:
-            raise Exception("Could not connect to the database.")
-
-    def query_dict(self, query_statement, query_parameters=None):
-        return self.query(query_statement, query_parameters, dictionary=True)
-
 
 def get_anunciosdb():
     return DatabaseWrapper(anuncios_db_parameters)