Enable mypy checking for dlstbx (#130)

Happy 2022

Co-authored-by: Richard Gildea <richard.gildea@diamond.ac.uk>

Author: Anthchirp

.pre-commit-config.yaml

@@ -32,3 +32,11 @@ repos:
   - id: check-merge-conflict
   - id: check-added-large-files
     args: ['--maxkb=200']
+
+# Type checking
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v0.910
+  hooks:
+  - id: mypy
+    files: 'src/.*\.py$'
+    additional_dependencies: ['types-pytz==2021.3.2', 'types-python-dateutil==2.8.3', 'types-setuptools==57.4.4', 'types-PyYAML==6.0.1', 'types-pycurl==7.44.1', 'types-requests==2.26.1', 'types-certifi==2021.10.8.0', 'types-pyOpenSSL==21.0.1']

src/dlstbx/cli/fix_cluster_jobs.py

@@ -103,8 +103,9 @@ def get_resubmission_id_for_job(cluster, jobid):
             if resubmission_id:
                 error_db[j] = resubmission_id
             else:
-                error_db[j] = str(uuid.uuid4())
-                cluster.qalter(j, ["-ac", "resubmission_id=" + error_db[j]])
+                unique_id = str(uuid.uuid4())
+                cluster.qalter(j, ["-ac", f"resubmission_id={unique_id}"])
+                error_db[j] = unique_id
             bar.update(1)
 
     removable_jobs = [j for j in error_db if error_db[j] in resubmission_db]

src/dlstbx/cli/status_monitor.py

@@ -3,6 +3,8 @@
 #   Starts a status monitor (what do you expect?)
 #
 
+from __future__ import annotations
+
 import argparse
 import curses
 import os
@@ -11,6 +13,7 @@
 import threading
 import time
 from pprint import pprint
+from typing import Any, Tuple
 
 import workflows
 import workflows.services
@@ -37,15 +40,15 @@ class Monitor:
     shutdown = False
     """Set to true to end the main loop and shut down the service monitor."""
 
-    cards = {}
+    cards: dict[int, Any] = {}
     """Register card shown for seen services"""
 
     border_chars = ()
     """Characters used for frame borders."""
     border_chars_text = ("|", "|", "=", "=", "/", "\\", "\\", "/")
     """Example alternative set of frame border characters."""
 
-    most_recent_version = {}
+    most_recent_version: dict[Any, Tuple[Any, Any]] = {}
     """Dictionary to hold software version information, so old versions can be highlighted."""
 
     def __init__(self, filters=None, transport=None, version=None, test=False):

src/dlstbx/health_checks/__init__.py

@@ -26,7 +26,7 @@ class CheckFunctionInterface(NamedTuple):
     name: str
 
 
-class Status(_Base):
+class Status(_Base):  # type: ignore
     __tablename__ = "infrastructure_status"
     __table_args__ = {"comment": "reports of DLS infrastructure"}
 

src/dlstbx/health_checks/quota.py

@@ -36,7 +36,7 @@ def _check_quota_for_user(
     entries = [entry for entry in _quotafile() if username in entry]
     limits = {"notice": 0.925, "warning": 0.95, "error": 0.975}
 
-    def report_level(value: float, limit: float) -> REPORT:
+    def report_level(value: float, limit: float) -> int:
         if value >= limits["error"] * limit:
             return REPORT.ERROR
         if value >= limits["warning"] * limit:
@@ -52,7 +52,7 @@ def report_level(value: float, limit: float) -> REPORT:
             Level=REPORT.ERROR,
             Message=f"Only {len(entries)} entries found in quota file for {username}",
         )
-    notices: list[tuple[REPORT, str]] = []
+    notices: list[tuple[int, str]] = []
 
     for entry in entries:
         tokens = entry.split()

src/dlstbx/services/images.py

@@ -3,7 +3,7 @@
 import os
 import re
 import time
-from typing import Any, Callable, Dict, NamedTuple
+from typing import Any, Callable, Dict, NamedTuple, Protocol
 
 import PIL.Image
 import pkg_resources
@@ -14,9 +14,14 @@
 logger = logging.getLogger("dlstbx.services.images")
 
 
+class _CallableParameter(Protocol):
+    def __call__(self, key: str, default: Any = ...) -> Any:
+        ...
+
+
 class PluginInterface(NamedTuple):
     rw: workflows.recipe.wrapper.RecipeWrapper
-    parameters: Callable[[str], Any]
+    parameters: _CallableParameter
     message: Dict[str, Any]
 
 
@@ -61,7 +66,7 @@ def image_call(self, rw, header, message):
         """Pass incoming message to the relevant plugin function."""
         command = rw.recipe_step.get("parameters", {}).get("image_command")
 
-        def parameters(key, default=None):
+        def parameters(key: str, default=None):
             if isinstance(message, dict) and message.get(key):
                 return message[key]
             return rw.recipe_step.get("parameters", {}).get(key, default)

src/dlstbx/services/trigger.py

@@ -1,10 +1,12 @@
+from __future__ import annotations
+
 import hashlib
 import logging
 import os
 import pathlib
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Dict, List, Literal, Mapping, Optional
+from typing import Any, Dict, List, Literal, Mapping, Optional
 
 import ispyb
 import pydantic
@@ -351,28 +353,28 @@ def trigger_dimple(
 
         dcid = parameters.dcid
 
-        pdb_files = self.get_linked_pdb_files_for_dcid(
+        pdb_files_or_codes = self.get_linked_pdb_files_for_dcid(
             session,
             dcid,
             parameters.pdb_tmpdir,
             user_pdb_dir=parameters.user_pdb_directory,
         )
 
-        if not pdb_files:
+        if not pdb_files_or_codes:
             self.log.info(
                 "Skipping dimple trigger: DCID %s has no associated PDB information"
                 % dcid
             )
             return {"success": True}
-        pdb_files = [str(p) for p in pdb_files]
+        pdb_files = [str(p) for p in pdb_files_or_codes]
         self.log.info("PDB files: %s", ", ".join(pdb_files))
 
         dc = (
             session.query(DataCollection)
             .filter(DataCollection.dataCollectionId == dcid)
             .one()
         )
-        dimple_parameters = {
+        dimple_parameters: dict[str, list[Any]] = {
             "data": [os.fspath(parameters.mtz)],
             "scaling_id": [parameters.scaling_id],
             "pdb": pdb_files,
@@ -766,17 +768,19 @@ def trigger_mrbump(
             self.log.info("Skipping mrbump trigger: sequence information not available")
             return {"success": True}
 
-        pdb_files = self.get_linked_pdb_files_for_dcid(
-            session,
-            dcid,
-            parameters.pdb_tmpdir,
-            user_pdb_dir=parameters.user_pdb_directory,
-            ignore_pdb_codes=True,
+        pdb_files = tuple(
+            self.get_linked_pdb_files_for_dcid(
+                session,
+                dcid,
+                parameters.pdb_tmpdir,
+                user_pdb_dir=parameters.user_pdb_directory,
+                ignore_pdb_codes=True,
+            )
         )
 
         jobids = []
 
-        for pdb_files in {(), tuple(pdb_files)}:
+        for pdb_files in {(), pdb_files}:
             jp = self.ispyb.mx_processing.get_job_params()
             jp["automatic"] = parameters.automatic
             jp["comments"] = parameters.comment
@@ -806,6 +810,7 @@ def trigger_mrbump(
                 self.log.debug(f"mrbump trigger: generated JobParameterID {jppid}")
 
             for pdb_file in pdb_files:
+                assert pdb_file.filepath is not None
                 filepath = pdb_file.filepath
                 if pdb_file.source == "AlphaFold":
                     trimmed = filepath.with_name(
@@ -1130,7 +1135,7 @@ def trigger_big_ep(
         class BigEPParams(pydantic.BaseModel):
             data: pathlib.Path
             scaled_unmerged_mtz: pathlib.Path
-            path_ext: Optional[str] = pydantic.Field(
+            path_ext: str = pydantic.Field(
                 default_factory=lambda: datetime.now().strftime("%Y%m%d_%H%M%S")
             )
 
@@ -1294,7 +1299,7 @@ def trigger_multiplex(
         status["ntry"] += 1
         self.log.debug(f"dcid={dcid}\nmessage_delay={message_delay}\n{status}")
 
-        multiplex_job_dcids = []
+        multiplex_job_dcids: list[set[int]] = []
         jobids = []
 
         for group in related_dcids:

src/dlstbx/test/services/test_trigger.py

@@ -765,7 +765,8 @@ def test_alphafold_not_triggered_for_industry_proposal(
     with caplog.at_level(logging.DEBUG):
         trigger.trigger(rw, {"some": "header"}, message)
     assert "Not triggering AlphaFold for protein_id" in caplog.text
-    send.assert_not_called()
+    send.assert_called_once_with({"result": None}, transaction=mocker.ANY)
+    t.send.assert_not_called()
 
 
 def test_invalid_params(db_session_factory, caplog):

src/dlstbx/util/certificate.py

@@ -15,9 +15,6 @@ def problems_with_certificate(hostname: str) -> Optional[str]:
         with closing(socket.socket()) as sock:
             sock.connect((hostname, 443))
             ctx = SSL.Context(SSL.SSLv23_METHOD)  # most compatible
-            ctx.check_hostname = False
-            ctx.verify_mode = SSL.VERIFY_NONE
-
             with closing(SSL.Connection(ctx, sock)) as sock_ssl:
                 sock_ssl.set_connect_state()
                 sock_ssl.set_tlsext_host_name(idna.encode(hostname))

src/dlstbx/util/colorstreamhandler.py

@@ -50,28 +50,23 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
+from __future__ import annotations
 
 import logging
 import platform
+from typing import Union
 
 
-class _AnsiColorStreamHandler(logging.StreamHandler):
-    DEFAULT = "\x1b[0m"
-    RED = "\x1b[31m"
-    GREEN = "\x1b[32m"
-    YELLOW = "\x1b[33m"
-    BLUE = "\x1b[34m"
-    CYAN = "\x1b[36m"
-    BOLD = "\x1b[1m"
-
-    CRITICAL = RED + BOLD
-    ERROR = RED
-    WARNING = YELLOW
-    INFO = GREEN
-    DEBUG = BLUE
+class _GenericColorStreamHandler(logging.StreamHandler):
+    DEFAULT: Union[int, str] = ""
+    CRITICAL: Union[int, str] = ""
+    ERROR: Union[int, str] = ""
+    WARNING: Union[int, str] = ""
+    INFO: Union[int, str] = ""
+    DEBUG: Union[int, str] = ""
 
     @classmethod
-    def _get_color(cls, level):
+    def _get_color(cls, level: int) -> Union[int, str]:
         if level >= logging.CRITICAL:
             return cls.CRITICAL
         elif level >= logging.ERROR:
@@ -85,16 +80,29 @@ def _get_color(cls, level):
         else:
             return cls.DEFAULT
 
-    def __init__(self, stream=None):
-        logging.StreamHandler.__init__(self, stream)
+
+class _AnsiColorStreamHandler(_GenericColorStreamHandler):
+    DEFAULT = "\x1b[0m"
+    RED = "\x1b[31m"
+    GREEN = "\x1b[32m"
+    YELLOW = "\x1b[33m"
+    BLUE = "\x1b[34m"
+    CYAN = "\x1b[36m"
+    BOLD = "\x1b[1m"
+
+    CRITICAL = RED + BOLD
+    ERROR = RED
+    WARNING = YELLOW
+    INFO = GREEN
+    DEBUG = BLUE
 
     def format(self, record):
         text = logging.StreamHandler.format(self, record)
         color = self._get_color(record.levelno)
         return color + text + self.DEFAULT
 
 
-class _WinColorStreamHandler(logging.StreamHandler):
+class _WinColorStreamHandler(_GenericColorStreamHandler):
     # wincon.h
     FOREGROUND_BLACK = 0x0000
     FOREGROUND_BLUE = 0x0001
@@ -126,28 +134,14 @@ class _WinColorStreamHandler(logging.StreamHandler):
     INFO = FOREGROUND_GREEN
     DEBUG = FOREGROUND_CYAN
 
-    @classmethod
-    def _get_color(cls, level):
-        if level >= logging.CRITICAL:
-            return cls.CRITICAL
-        elif level >= logging.ERROR:
-            return cls.ERROR
-        elif level >= logging.WARNING:
-            return cls.WARNING
-        elif level >= logging.INFO:
-            return cls.INFO
-        elif level >= logging.DEBUG:
-            return cls.DEBUG
-        else:
-            return cls.DEFAULT
-
     def _set_color(self, code):
         import ctypes
 
         ctypes.windll.kernel32.SetConsoleTextAttribute(self._outhdl, code)
 
     def __init__(self, stream=None):
-        logging.StreamHandler.__init__(self, stream)
+        super().__init__(stream)
+
         # get file handle for the stream
         import ctypes
         import ctypes.util
@@ -167,6 +161,7 @@ def emit(self, record):
 
 
 # select ColorStreamHandler based on platform
+ColorStreamHandler: type[_GenericColorStreamHandler]
 if platform.system() == "Windows":
     ColorStreamHandler = _WinColorStreamHandler
 else: