Use correct classes for type-hints.

1 files changed, 14 insertions, 14 deletions

diff --git a/uploader/phenotypes/models.py b/uploader/phenotypes/models.py
index a22497c..5acdf60 100644
--- a/uploader/phenotypes/models.py
+++ b/uploader/phenotypes/models.py
@@ -6,8 +6,8 @@ from functools import reduce
 from datetime import datetime
 from typing import Union, Optional, Iterable
 
-import MySQLdb as mdb
-from MySQLdb.cursors import Cursor, DictCursor
+from MySQLdb.connections import Connection
+from MySQLdb.cursors import Cursor, DictCursor, BaseCursor
 
 from gn_libs.mysqldb import debug_query
 
@@ -27,7 +27,7 @@ __PHENO_DATA_TABLES__ = {
 
 
 def datasets_by_population(
-        conn: mdb.Connection,
+        conn: Connection,
         species_id: int,
         population_id: int
 ) -> tuple[dict, ...]:
@@ -42,7 +42,7 @@ def datasets_by_population(
         return tuple(dict(row) for row in cursor.fetchall())
 
 
-def dataset_by_id(conn: mdb.Connection,
+def dataset_by_id(conn: Connection,
                   species_id: int,
                   population_id: int,
                   dataset_id: int) -> dict:
@@ -57,7 +57,7 @@ def dataset_by_id(conn: mdb.Connection,
         return dict(cursor.fetchone())
 
 
-def phenotypes_count(conn: mdb.Connection,
+def phenotypes_count(conn: Connection,
                      population_id: int,
                      dataset_id: int) -> int:
     """Count the number of phenotypes in the dataset."""
@@ -85,7 +85,7 @@ def phenotype_publication_data(conn, phenotype_id) -> Optional[dict]:
         return dict(res)
 
 
-def dataset_phenotypes(conn: mdb.Connection,
+def dataset_phenotypes(conn: Connection,
                        population_id: int,
                        dataset_id: int,
                        offset: int = 0,
@@ -105,7 +105,7 @@ def dataset_phenotypes(conn: mdb.Connection,
         return tuple(dict(row) for row in cursor.fetchall())
 
 
-def __phenotype_se__(cursor: Cursor, xref_id, dataids_and_strainids):
+def __phenotype_se__(cursor: BaseCursor, xref_id, dataids_and_strainids):
     """Fetch standard-error values (if they exist) for a phenotype."""
     paramstr = ", ".join(["(%s, %s)"] * len(dataids_and_strainids))
     flat = tuple(item for sublist in dataids_and_strainids for item in sublist)
@@ -187,7 +187,7 @@ def __merge_pheno_data_and_se__(data, sedata) -> dict:
 
 
 def phenotype_by_id(
-        conn: mdb.Connection,
+        conn: Connection,
         species_id: int,
         population_id: int,
         dataset_id: int,
@@ -225,7 +225,7 @@ def phenotype_by_id(
     return None
 
 
-def phenotypes_data(conn: mdb.Connection,
+def phenotypes_data(conn: Connection,
                     population_id: int,
                     dataset_id: int,
                     offset: int = 0,
@@ -249,7 +249,7 @@ def phenotypes_data(conn: mdb.Connection,
 
 
 def phenotypes_vector_data(# pylint: disable=[too-many-arguments, too-many-positional-arguments]
-        conn: mdb.Connection,
+        conn: Connection,
         species_id: int,
         population_id: int,
         xref_ids: tuple[int, ...] = tuple(),
@@ -301,7 +301,7 @@ def phenotypes_vector_data(# pylint: disable=[too-many-arguments, too-many-posit
         return reduce(__organise__, cursor.fetchall(), {})
 
 
-def save_new_dataset(cursor: Cursor,
+def save_new_dataset(cursor: BaseCursor,
                      population_id: int,
                      dataset_name: str,
                      dataset_fullname: str,
@@ -346,7 +346,7 @@ def __pre_process_phenotype_data__(row):
 
 
 def create_new_phenotypes(# pylint: disable=[too-many-locals]
-        conn: mdb.Connection,
+        conn: Connection,
         population_id: int,
         publication_id: int,
         phenotypes: Iterable[dict]
@@ -474,7 +474,7 @@ def create_new_phenotypes(# pylint: disable=[too-many-locals]
 
 
 def save_phenotypes_data(
-        conn: mdb.Connection,
+        conn: Connection,
         table: str,
         data: Iterable[dict]
 ) -> int:
@@ -504,7 +504,7 @@ def save_phenotypes_data(
 
 
 def quick_save_phenotypes_data(
-        conn: mdb.Connection,
+        conn: Connection,
         table: str,
         dataitems: Iterable[dict],
         tmpdir: Path