YDA-5992: add type annotations to utils

.github/workflows/python.yml

@@ -19,7 +19,8 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install flake8==6.0.0 flake8-import-order==0.18.2 darglint==1.8.1 codespell mypy types-requests types-python-dateutil
+          python -m pip install flake8==6.0.0 flake8-import-order==0.18.2 darglint==1.8.1 codespell
+          python -m pip install mypy types-requests types-python-dateutil types-redis
 
       - name: Lint with flake8
         run: |

setup.cfg

@@ -17,6 +17,6 @@ check_untyped_defs = False
 disallow_any_generics = False
 disallow_incomplete_defs = True
 disallow_untyped_calls = False
-disallow_untyped_defs = True
+disallow_untyped_defs = False
 show_error_codes = True
 show_error_context = True

util/api.py

@@ -11,28 +11,29 @@
 import traceback
 import zlib
 from collections import OrderedDict
+from typing import Callable, Dict
 
+import error
 import jsonutil
 import log
 import rule
 from config import config
-from error import *
 
 
 class Result:
     """API result."""
 
-    def __init__(self, data=None, status='ok', info=None, debug_info=None):
+    def __init__(self, data: Dict | None = None, status: str = 'ok', info: str | None = None, debug_info: str | None = None) -> None:
         self.status      = status
         self.status_info = info
         self.data        = data
         self.debug_info  = debug_info
 
     @staticmethod
-    def ok(**xs):
+    def ok(**xs: int) -> object:
         return Result(**xs)
 
-    def as_dict(self):
+    def as_dict(self) -> OrderedDict:
         if config.environment == 'development':
             # Emit debug information in dev.
             # This may contain stack traces, exception texts, timing info,
@@ -46,29 +47,29 @@ def as_dict(self):
                                 ('status_info', self.status_info),
                                 ('data',        self.data)])
 
-    def __bool__(self):
+    def __bool__(self) -> bool:
         return self.status == 'ok'
     __nonzero__ = __bool__
 
 
-class Error(Result, UUError):
+class Error(Result, error.UUError):
     """Error with descriptive (user-readable) message.
 
     Returned/raised by API functions to produce informative error output.
     """
-    def __init__(self, name, info, debug_info=None, data=None):
+    def __init__(self, name: str, info: str, debug_info: str | None = None, data: str | None = None) -> None:
         self.name = name
         self.info = info
         self.debug_info = debug_info
 
         Result.__init__(self, data, 'error_' + name, info, debug_info)
-        UUError.__init__(self, 'error_' + name)
+        error.UUError.__init__(self, 'error_' + name)
 
-    def __str__(self):
+    def __str__(self) -> str:
         return '{}: {}'.format(self.name, self.info)
 
 
-def _api(f):
+def _api(f: Callable) -> Callable:
     """Turn a Python function into a basic API function.
 
     By itself, this wrapper is not very useful, as the resulting function is
@@ -107,7 +108,7 @@ def _api(f):
     # If the function accepts **kwargs, we do not forbid extra arguments.
     allow_extra = a_kw is not None
 
-    def wrapper(ctx, inp):
+    def wrapper(ctx: rule.Context, inp: str) -> Dict:
         """A function that receives a JSON string and calls a wrapped function with unpacked arguments.
 
         :param ctx: Combined type of a callback and rei struct
@@ -119,10 +120,10 @@ def wrapper(ctx, inp):
         :returns: Result of the JSON API call
         """
         # Result shorthands.
-        def error_internal(debug_info=None):
+        def error_internal(debug_info: str | None = None) -> Error:
             return Error('internal', 'An internal error occurred', debug_info=debug_info)
 
-        def bad_request(debug_info=None):
+        def bad_request(debug_info: str | None = None) -> Error:
             return Error('badrequest', 'An internal error occurred', debug_info=debug_info)
 
         # Input is base64 encoded and compressed to reduce size (max rule length in iRODS is 20KB)
@@ -195,7 +196,7 @@ def bad_request(debug_info=None):
     return wrapper
 
 
-def make():
+def make() -> Callable:
     """Create API functions callable as iRODS rules.
 
     This translate between a Python calling convention and the iRODS rule
@@ -221,7 +222,7 @@ def api_ping(ctx, foo):
 
     :returns: API function callable as iRODS rules
     """
-    def deco(f):
+    def deco(f: Callable) -> Callable:
         # The "base" API function, that does handling of arguments and errors.
         base = _api(f)
 

util/arb_data_manager.py

@@ -6,53 +6,55 @@
 __copyright__ = 'Copyright (c) 2019-2024, Utrecht University'
 __license__   = 'GPLv3, see LICENSE'
 
+from typing import TYPE_CHECKING
+
 import genquery
 
 import cached_data_manager
 import constants
 import log
 import msi
+if TYPE_CHECKING:
+    import rule
 
 
 class ARBDataManager(cached_data_manager.CachedDataManager):
     AVU_NAME = "yoda::arb"
 
-    def get(self, ctx, keyname):
-        """Retrieves data from the cache if possible, otherwise retrieves
-           the original.
+    def get(self, ctx: 'rule.Context', keyname: str) -> str:
+        """Retrieves data from the cache if possible, otherwise retrieves the original.
 
-           :param ctx:     Combined type of a callback and rei struct
-           :param keyname: name of the key
+        :param ctx:     Combined type of a callback and rei struct
+        :param keyname: Name of the key
 
-           :returns:       data for this key (arb_status)
+        :returns: Data for this key (arb_status)
         """
         value = super().get(ctx, keyname)
         return constants.arb_status[value.decode("utf-8")]
 
-    def put(self, ctx, keyname, data):
+    def put(self, ctx: 'rule.Context', keyname: str, data: str) -> None:
         """Update both the original value and cached value (if cache is not available, it is not updated)
 
-           :param ctx:     Combined type of a callback and rei struct
-           :param keyname: name of the key
-           :param data: data for this key (arb_status)
+        :param ctx:     Combined type of a callback and rei struct
+        :param keyname: Name of the key
+        :param data:    Data for this key (arb_status)
         """
         super().put(ctx, keyname, data.value)
 
-    def _get_context_string(self):
-        """ :returns: a string that identifies the particular type of data manager
+    def _get_context_string(self) -> str:
+        """Returns a string that identifies the particular type of data manager.
 
-           :returns: context string for this type of data manager
+        :returns: context string for this type of data manager
         """
         return "arb"
 
-    def _get_original_data(self, ctx, keyname):
-        """This function is called when data needs to be retrieved from the original
-           (non-cached) location.
+    def _get_original_data(self, ctx: 'rule.Context', keyname: str) -> str:
+        """This function is called when data needs to be retrieved from the original (non-cached) location.
 
-           :param ctx:     Combined type of a callback and rei struct
-           :param keyname: name of the key
+        :param ctx:     Combined type of a callback and rei struct
+        :param keyname: Name of the key
 
-           :returns:       Original data for this key
+        :returns: Original data for this key
         """
         arb_data = list(genquery.row_iterator(
             "META_RESC_ATTR_VALUE",
@@ -68,21 +70,19 @@ def _get_original_data(self, ctx, keyname):
             log.write(ctx, "WARNING: multiple ARB AVUs present for resource '{}'. ARB will ignore it.".format(keyname))
             return constants.arb_status.IGNORE.value
 
-    def _put_original_data(self, ctx, keyname, data):
-        """This function is called when data needs to be updated in the original
-           (non-cached) location.
+    def _put_original_data(self, ctx: 'rule.Context', keyname: str, data: str) -> None:
+        """This function is called when data needs to be updated in the original (non-cached) location.
 
-           :param ctx:     Combined type of a callback and rei struct
-           :param keyname: name of the key
-           :param data:    Data for this key
+        :param ctx:     Combined type of a callback and rei struct
+        :param keyname: Name of the key
+        :param data:    Data for this key
         """
         msi.mod_avu_metadata(ctx, "-r", keyname, "set", self.AVU_NAME, data, "")
 
-    def _should_populate_cache_on_get(self):
-        """This function controls whether the manager populates the cache
-           after retrieving original data.
+    def _should_populate_cache_on_get(self) -> bool:
+        """This function controls whether the manager populates the cache after retrieving original data.
 
-           :returns: Boolean value that states whether the cache should be populated when original data
-                     is retrieved.
+        :returns: Boolean value that states whether the cache should be populated when original data
+                  is retrieved.
         """
         return True