Add type hints for qserv kraken

Author: fritzm

pyproject.toml

@@ -51,7 +51,6 @@ exclude = [
     "etc/*",
     "extern/*",
     "python/lsst/utils",
-    "python/lsst/qserv/testing",
     "python/lsst/qserv/.*/tests",
     "src/*",
 ]

python/lsst/qserv/testing/__init__.py

@@ -1,3 +1,6 @@
 from .config import Config, QueryFactory
 
-__all__ = [Config, QueryFactory]
+__all__ = [
+    "Config",
+    "QueryFactory",
+]

python/lsst/qserv/testing/config.py

@@ -5,6 +5,8 @@
 import contextlib
 import logging
 import random
+from collections.abc import Generator
+from typing import Self, TextIO
 
 import numpy as np
 import yaml
@@ -13,15 +15,15 @@
 
 
 @contextlib.contextmanager
-def _make_file(path_or_file):
+def _make_file(path_or_file: str | TextIO) -> Generator[TextIO, None, None]:
     """Context manager that makes a file out of argument.
 
     Parameters
     ----------
     path_or_file : `str` or file object
         Path name for a file or a file object.
     """
-    if hasattr(path_or_file, "read"):
+    if isinstance(path_or_file, TextIO):
         yield path_or_file
     else:
         with open(path_or_file) as file:
@@ -37,11 +39,11 @@ class _ValueRandomUniform:
         Range for generated numbers.
     """
 
-    def __init__(self, min, max):
+    def __init__(self: Self, min: float, max: float) -> None:
         self._min = float(min)
         self._max = float(max)
 
-    def __call__(self):
+    def __call__(self) -> float:
         return random.uniform(self._min, self._max)
 
 
@@ -54,11 +56,11 @@ class _ValueRandomUniformInt:
         Range for generated numbers.
     """
 
-    def __init__(self, min, max):
+    def __init__(self: Self, min: int, max: int) -> None:
         self._min = float(min)
         self._max = float(max)
 
-    def __call__(self):
+    def __call__(self) -> int:
         return int(random.uniform(self._min, self._max))
 
 
@@ -73,7 +75,9 @@ class _ValueIntFromFile:
         One of "random" or "sequential".
     """
 
-    def __init__(self, path, mode="random"):
+    _array: np.ndarray | list[int]
+
+    def __init__(self: Self, path: str, mode: str = "random") -> None:
         # read all numbers from file as integers
         if path == "/dev/null":
             # for testing only
@@ -84,7 +88,7 @@ def __init__(self, path, mode="random"):
         assert mode in ("random", "sequential")
         self._seq = 0
 
-    def __call__(self):
+    def __call__(self) -> int:
         if self._mode == "random":
             return random.choice(self._array)
         else:
@@ -108,12 +112,12 @@ class QueryFactory:
         with a description of how to generate variable value.
     """
 
-    def __init__(self, txt, variables=None):
+    def __init__(self: Self, txt: str, variables: dict[str, dict] | None = None) -> None:
         self._txt = txt
         self._vars = {}
         if variables is not None:
             for var, config in variables.items():
-                generator = None
+                generator: _ValueRandomUniform | _ValueRandomUniformInt | _ValueIntFromFile | None = None
                 if "distribution" in config:
                     if config["distribution"] == "uniform":
                         min = config.get("min", 0.0)
@@ -131,7 +135,7 @@ def __init__(self, txt, variables=None):
                     raise ValueError(f"Cannot parse variable configuration {var} = {config}")
                 self._vars[var] = generator
 
-    def query(self):
+    def query(self: Self) -> str:
         """Return next query to execute.
 
         Returns
@@ -157,7 +161,10 @@ class Config:
         List of dictionaries, cannot be empty.
     """
 
-    def __init__(self, configs):
+    _config: dict
+    _queries: dict
+
+    def __init__(self: Self, configs: list[dict]) -> None:
         if not configs:
             raise ValueError("empty configurations list")
 
@@ -190,7 +197,7 @@ def __init__(self, configs):
                     raise ValueError(f"Unexpected query configuration: {qkey}: {qcfg}")
 
     @classmethod
-    def from_yaml(cls, config_files):
+    def from_yaml(cls: type[Self], config_files: list[str | TextIO]) -> Self:
         """Make configuration from bunch of YAML files
 
         Parameters
@@ -210,7 +217,7 @@ def from_yaml(cls, config_files):
                 configs.append(yaml.load(file, Loader=yaml.SafeLoader))
         return cls(configs)
 
-    def to_yaml(self):
+    def to_yaml(self: Self) -> str:
         """Convert current config to YAML string.
 
         Returns
@@ -220,7 +227,7 @@ def to_yaml(self):
         """
         return yaml.dump(self._config)
 
-    def classes(self):
+    def classes(self: Self) -> set[str]:
         """Return set of classes defined in configuration.
 
         Returns
@@ -230,7 +237,7 @@ def classes(self):
         """
         return self._classes
 
-    def queries(self, q_class):
+    def queries(self: Self, q_class: str) -> dict[str, QueryFactory]:
         """Return queries for given class.
 
         Parameters
@@ -246,7 +253,7 @@ def queries(self, q_class):
         """
         return self._queries[q_class]
 
-    def concurrent_queries(self, q_class):
+    def concurrent_queries(self: Self, q_class: str) -> int:
         """Return number of concurrent queries for given class.
 
         Parameters
@@ -261,7 +268,7 @@ def concurrent_queries(self, q_class):
         """
         return self._config["queryClasses"][q_class]["concurrentQueries"]
 
-    def max_rate(self, q_class):
+    def max_rate(self: Self, q_class: str) -> float:
         """Return maximum rate for given class.
 
         Parameters
@@ -276,7 +283,7 @@ def max_rate(self, q_class):
         """
         return self._config["queryClasses"][q_class].get("maxRate")
 
-    def arraysize(self, q_class):
+    def arraysize(self: Self, q_class: str) -> int:
         """Return array size for fetchmany().
 
         Parameters
@@ -291,7 +298,7 @@ def arraysize(self, q_class):
         """
         return self._config["queryClasses"][q_class].get("arraysize")
 
-    def split(self, n_workers, i_worker):
+    def split(self: Self, n_workers: int, i_worker: int) -> Self:
         """Divide configuration (or its workload) between number of workers.
 
         If we want to run test with multiple workers we need to divide work
@@ -328,7 +335,7 @@ def split(self, n_workers, i_worker):
         return self.__class__([self._config, dict(queryClasses=overrides)])
 
     @staticmethod
-    def _merge(config1, config2):
+    def _merge(config1: dict, config2: dict) -> dict:
         """Merge two config objects, return result.
 
         If configuration present in both then second one overrides first.

python/lsst/qserv/testing/main.py

@@ -8,18 +8,22 @@
 import functools
 import logging
 import os
+from collections.abc import Callable
 
-import MySQLdb
+from lsst.qserv.testing import mock_db
+
+import mysql.connector
+from mysql.connector.abstracts import MySQLConnectionAbstract
+from mysql.connector.pooling import PooledMySQLConnection
 
-from . import mock_db
 from .config import Config
 from .monitor import InfluxDBFileMonitor, LogMonitor
 from .runner_mgr import RunnerManager
 
 _LOG = logging.getLogger(__name__)
 
 
-def _log_config(level, slot):
+def _log_config(level: int, slot: int | None) -> None:
     levels = {0: logging.WARNING, 1: logging.INFO, 2: logging.DEBUG}
     if slot is None:
         simple_format = "%(asctime)s %(levelname)7s %(name)s -- %(message)s"
@@ -29,7 +33,7 @@ def _log_config(level, slot):
     logging.basicConfig(format=simple_format, level=level)
 
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser(description="Test harness to generate load for QServ")
 
     parser.add_argument(
@@ -141,9 +145,9 @@ def main():
         if num_slots is None or slot is None:
             parser.error("cannot determine slurm configuration, envvar is not set")
             return 2
-        num_slots = int(num_slots)
-        if num_slots > 1:
-            args.num_slots = num_slots
+        num_slots_int = int(num_slots)
+        if num_slots_int > 1:
+            args.num_slots = num_slots_int
             args.slot = int(slot)
 
     _log_config(args.verbose, args.slot)
@@ -157,16 +161,22 @@ def main():
         print(cfg.to_yaml())
 
     # connection factory
+    conn_factory: Callable[..., PooledMySQLConnection | MySQLConnectionAbstract | mock_db.MockConnection]
     if args.dummy_db:
         conn_factory = mock_db.connect
     else:
         conn_factory = functools.partial(
-            MySQLdb.connect, host=args.host, port=args.port, user=args.user, passwd=args.password, db=args.db
+            mysql.connector.connect,
+            host=args.host,
+            port=args.port,
+            user=args.user,
+            passwd=args.password,
+            db=args.db,
         )
 
     # monitor
     tags = None if args.slot is None else {"slot": args.slot}
-    monitor = None
+    monitor: LogMonitor | InfluxDBFileMonitor | None = None
     if args.monitor == "log":
         monitor = LogMonitor(logging.getLogger("metrics"), tags=tags)
     elif args.monitor == "influxdb-file":

python/lsst/qserv/testing/mock_db.py

@@ -4,13 +4,16 @@
 import re
 import time
 from collections import namedtuple
+from typing import Any, Self
 
 _LOG = logging.getLogger(__name__)
 
 _num_re = re.compile(r"\d+")
 
 
-ColDesriptor = namedtuple("ColDesriptor", "name type_code display_size internal_size precision scale null_ok")
+ColDescriptor = namedtuple(
+    "ColDescriptor", "name type_code display_size internal_size precision scale null_ok"
+)
 
 
 class MockCursor:
@@ -20,13 +23,16 @@ class MockCursor:
     """
 
     arraysize = 1
+    _query: str | None
+    n_rows: int
+    rows: list[tuple[int, str]]
 
-    def __init__(self):
+    def __init__(self: Self) -> None:
         self._query = None
         self.n_rows = 0
         self.rows = []
 
-    def execute(self, query):
+    def execute(self: Self, query: str) -> None:
         _LOG.debug("executing query: %s", query)
         self._query = query
         self.n_rows = 2
@@ -40,30 +46,30 @@ def execute(self, query):
         # spend at least few milliseconds in query
         time.sleep(0.01)
 
-    def fetchall(self):
+    def fetchall(self: Self) -> list[tuple[int, str]]:
         rows = self.rows
         self.rows = []
         return rows
 
-    def fetchmany(self, arraysize=None):
+    def fetchmany(self: Self, arraysize: int | None = None) -> list[tuple[int, str]]:
         if arraysize is None:
             arraysize = self.arraysize
         rows = self.rows[:arraysize]
         self.rows = self.rows[arraysize:]
         return rows
 
     @property
-    def rowcount(self):
+    def rowcount(self: Self) -> int:
         return self.n_rows
 
     @property
-    def description(self):
+    def description(self: Self) -> list[ColDescriptor]:
         # some randome codes
         return [
-            ColDesriptor(
+            ColDescriptor(
                 name="ID", type_code=1, display_size=10, internal_size=4, precision=0, scale=1, null_ok=False
             ),
-            ColDesriptor(
+            ColDescriptor(
                 name="name",
                 type_code=15,
                 display_size=32,
@@ -76,11 +82,11 @@ def description(self):
 
 
 class MockConnection:
-    def cursor(self):
+    def cursor(self: Self) -> MockCursor:
         return MockCursor()
 
 
-def connect(*args, **kwargs):
+def connect(*args: Any, **kwargs: Any) -> MockConnection:
     """Can take any parameters so it can be used as replacement for
     MySQLdb.connect() method.
     """