Add mypy type hints for utils

src/drunc/utils/__init__.py

@@ -1,3 +1,5 @@
+"""drunc utilities module."""
+
 from drunc.utils.utils import get_logger
 
 # Initialise utils logger with Rich handler

src/drunc/utils/configuration.py

@@ -1,6 +1,9 @@
+"""Configuration utilities for DRUNC."""
+
 import json
 import os
 from enum import Enum
+from typing import Protocol, cast
 
 import conffwk
 
@@ -9,6 +12,8 @@
 
 
 class ConfTypes(Enum):
+    """Enumeration of supported configuration types."""
+
     Unknown = 0
 
     # End product
@@ -21,6 +26,17 @@ class ConfTypes(Enum):
 
 
 def CLI_to_ConfTypes(scheme: str) -> ConfTypes:
+    """Convert a CLI scheme string to a ConfTypes enum.
+
+    Args:
+        scheme: The scheme string ("file", "oksconflibs", or "").
+
+    Returns:
+        ConfTypes: The corresponding configuration type.
+
+    Raises:
+        DruncSetupException: If the scheme is not recognized.
+    """
     match scheme:
         case "file":
             return ConfTypes.JsonFileName
@@ -31,43 +47,106 @@ def CLI_to_ConfTypes(scheme: str) -> ConfTypes:
 
 
 def parse_conf_url(url: str) -> tuple[str, ConfTypes]:
+    """Parse a configuration URL into scheme and type.
+
+    Args:
+        url: The configuration URL (format: "scheme:filename").
+
+    Returns:
+        tuple[str, ConfTypes]: A tuple of (url, conf_type).
+    """
     scheme, filename = url.split(":")
     t = CLI_to_ConfTypes(scheme)
     return url, t
 
 
 class ConfigurationNotFound(DruncSetupException):
-    def __init__(self, requested_path):
+    """Exception raised when configuration is not found."""
+
+    def __init__(self, requested_path: str) -> None:
+        """Initialize the ConfigurationNotFound exception.
+
+        Args:
+            requested_path: The path to the configuration that was not found.
+        """
         super().__init__(
             f"The configuration '{requested_path}' is not in $DUNEDAQ_DB_PATH, perhaps you forgot to 'dbt-workarea-env && dbt-build'?"
         )
 
 
 class ConfTypeNotSupported(DruncSetupException):
-    def __init__(self, conf_type: ConfTypes, class_name: str):
+    """Exception raised when a configuration type is not supported."""
+
+    def __init__(self, conf_type: ConfTypes, class_name: str) -> None:
+        """Initialize the ConfTypeNotSupported exception.
+
+        Args:
+            conf_type: The configuration type that is not supported.
+            class_name: The name of the class where this type is not supported.
+        """
         if not isinstance(class_name, str):
             class_name = class_name.__class__.__name__
         message = f"'{conf_type}' is not supported by '{class_name}'"
         super().__init__(message)
 
 
 class OKSKey:
-    def __init__(self, schema_file: str, class_name: str, obj_uid: str, session: str):
+    """Key information for accessing OKS configuration objects."""
+
+    def __init__(
+        self, schema_file: str, class_name: str, obj_uid: str, session: str
+    ) -> None:
+        """Initialize an OKSKey.
+
+        Args:
+            schema_file: The OKS schema file path.
+            class_name: The class name in the OKS schema.
+            obj_uid: The unique identifier for the object.
+            session: The session name.
+        """
         self.schema_file = schema_file
         self.class_name = class_name
         self.obj_uid = obj_uid
         self.session = session
 
 
+class _DataTypeName(Protocol):
+    _name_: str
+
+
+class _ConfigurationData(Protocol):
+    type: _DataTypeName
+    broadcaster: object
+    authoriser: object
+
+
 class ConfHandler:
+    """Handler for loading and parsing DRUNC configurations.
+
+    Supports multiple configuration types including JSON files, Protobuf messages,
+    and OKS.
+    """
+
     def __init__(
         self,
-        data=None,
-        type=ConfTypes.PyObject,
-        oks_key: OKSKey = None,
-        *args,
-        **kwargs,
-    ):
+        data: object = None,
+        type: ConfTypes = ConfTypes.PyObject,
+        oks_key: OKSKey | None = None,
+        *args: object,
+        **kwargs: object,
+    ) -> None:
+        """Initialize a ConfHandler.
+
+        Args:
+            data: The configuration data. Defaults to None.
+            type: The configuration type. Defaults to PyObject.
+            oks_key: OKS key if using OKS configuration. Defaults to None.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Raises:
+            DruncSetupException: If OKS type is used without an OKS key.
+        """
         self.class_name = self.__class__.__name__
         self.log = get_logger("utils." + self.class_name)
         self.initial_type = type
@@ -84,26 +163,52 @@ def __init__(
         self.oks_key = oks_key
         self.validate_and_parse_configuration_location(*args, **kwargs)
 
-    def get_data(self):
+    def get_data(self) -> object:
+        """Get the configuration data.
+
+        Returns:
+            Any: The stored configuration data.
+        """
         return self.data
 
-    def get_data_type_name(self):
-        return self.get_data().type._name_
+    def get_data_type_name(self) -> str:
+        """Get the type name of the configuration data.
 
-    def get_data_broadcaster(self):
-        return self.get_data().broadcaster
+        Returns:
+            str: The name of the data type.
+        """
+        return str(cast(_ConfigurationData, self.get_data()).type._name_)
 
-    def get_data_authoriser(self):
-        return self.get_data().authoriser
+    def get_data_broadcaster(self) -> object:
+        """Get the broadcaster from the configuration data.
 
-    def copy_oks_key(self):
+        Returns:
+            Any: The broadcaster object.
+        """
+        return cast(_ConfigurationData, self.get_data()).broadcaster
+
+    def get_data_authoriser(self) -> object:
+        """Get the authoriser from the configuration data.
+
+        Returns:
+            Any: The authoriser object.
+        """
+        return cast(_ConfigurationData, self.get_data()).authoriser
+
+    def copy_oks_key(self) -> OKSKey | None:
+        """Get a copy of the OKS key if one exists.
+
+        Returns:
+            OKSKey | None: The OKS key, or None if not using OKS configuration.
+        """
         return self.oks_key
 
-    def _parse_oks_file(self, oks_path):
+    def _parse_oks_file(self, oks_path: str) -> object:
         try:
             self.oks_path = oks_path
             self.log.debug(f"Using {self.oks_path} to configure")
             self.db = conffwk.Configuration(self.oks_path)
+            assert self.oks_key is not None, "OKS key is required for OKS configuration"
             return self.db.get_dal(
                 class_name=self.oks_key.class_name, uid=self.oks_key.obj_uid
             )
@@ -118,25 +223,35 @@ def _parse_oks_file(self, oks_path):
                 "OKS params where not passed to this ConfigurationHandler, cannot parse OKS configurations"
             ) from e
 
-    def _post_process_oks(self):
+    def _post_process_oks(self, *args: object, **kwargs: object) -> None:
         pass
 
-    def _parse_pbany(self, pbany_data):
-        raise ConfTypeNotSupported(ConfTypes.ProtobufAny, self)
+    def _parse_pbany(self, pbany_data: object) -> object:
+        raise ConfTypeNotSupported(ConfTypes.ProtobufAny, self.class_name)
+
+    def _parse_dict(self, data: dict[str, object]) -> object:
+        raise ConfTypeNotSupported(ConfTypes.JsonFileName, self.class_name)
+
+    def validate_and_parse_configuration_location(
+        self, *args: object, **kwargs: object
+    ) -> None:
+        """Validate and parse the configuration from the provided location.
 
-    def _parse_dict(self, data):
-        raise ConfTypeNotSupported(ConfTypes.JsonFileName, self)
+        Supports JsonFileName, OKSFileName, and PyObject types.
 
-    def validate_and_parse_configuration_location(self, *args, **kwargs):
+        Args:
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
         match self.initial_type:
             case ConfTypes.PyObject:
                 self.data = self.initial_data
                 self.type = self.initial_type
                 self._post_process_oks(*args, **kwargs)
 
             case ConfTypes.JsonFileName:
-                resolved = expand_path(self.initial_data, True)
-                if not os.path.exists(expand_path(self.initial_data)):
+                resolved = expand_path(cast(str, self.initial_data), True)
+                if not os.path.exists(expand_path(cast(str, self.initial_data))):
                     raise DruncSetupException(
                         f"Location {resolved} ({self.initial_data}) is empty!"
                     )
@@ -148,7 +263,7 @@ def validate_and_parse_configuration_location(self, *args, **kwargs):
                     self._post_process_oks(*args, **kwargs)
 
             case ConfTypes.OKSFileName:
-                self.data = self._parse_oks_file(self.initial_data)
+                self.data = self._parse_oks_file(cast(str, self.initial_data))
                 self.type = ConfTypes.PyObject
                 self._post_process_oks(*args, **kwargs)