Add types for connection objects.

1 files changed, 9 insertions, 2 deletions

diff --git a/uploader/publications/models.py b/uploader/publications/models.py
index e7f197a..f36c800 100644
--- a/uploader/publications/models.py
+++ b/uploader/publications/models.py
@@ -1,7 +1,11 @@
 """Module to handle persistence and retrieval of publication to/from MariaDB"""
 
+from gn_libs.mysqldb import Connection
+
 def fetch_phenotype_publications(
-        conn, ids: tuple[tuple[int, int], ...]) -> tuple[dict, ...]:
+        conn: Connection,
+        ids: tuple[tuple[int, int], ...]
+) -> tuple[dict, ...]:
     """Fetch publication from database by ID."""
     paramstr = ",".join(["(%s, %s)"] * len(ids))
     query = (
@@ -15,7 +19,10 @@ def fetch_phenotype_publications(
         return tuple(dict(row) for row in cursor.fetchall())
 
 
-def create_new_publications(conn, publications) -> dict:
+def create_new_publications(
+        conn: Connection,
+        publications: tuple[dict, ...]
+) -> tuple[dict, ...]:
     if len(publications) > 0:
         with conn.cursor(cursorclass=DictCursor) as cursor:
             cursor.executemany(