Add utility decorator to log time spent in function calls

src/zino/scheduler.py

@@ -12,6 +12,7 @@
 from zino.config.models import DEFAULT_INTERVAL_MINUTES, PollDevice
 from zino.config.polldevs import InvalidConfiguration, read_polldevs
 from zino.tasks import run_all_tasks
+from zino.utils import log_time_spent
 
 _log = logging.getLogger(__name__)
 _scheduler = None
@@ -38,6 +39,7 @@ def get_scheduler() -> AsyncIOScheduler:
     return _scheduler
 
 
+@log_time_spent()
 def load_polldevs(polldevs_conf: str) -> Tuple[Set, Set]:
     """Loads pollfile into process state.
 

src/zino/state.py

@@ -13,6 +13,7 @@
 from zino.flaps import FlappingStates
 from zino.planned_maintenance import PlannedMaintenances
 from zino.statemodels import DeviceStates
+from zino.utils import log_time_spent
 
 _log = logging.getLogger(__name__)
 
@@ -34,13 +35,15 @@ class ZinoState(BaseModel):
     planned_maintenances: PlannedMaintenances = Field(default_factory=PlannedMaintenances)
     flapping: FlappingStates = Field(default_factory=FlappingStates)
 
+    @log_time_spent()
     def dump_state_to_file(self, filename: str):
         """Dumps the full state to a file in JSON format"""
         _log.debug("dumping state to %s", filename)
         with open(filename, "w") as statefile:
             statefile.write(self.model_dump_json(exclude_none=True, indent=2))
 
     @classmethod
+    @log_time_spent()
     def load_state_from_file(cls, filename: str) -> Optional["ZinoState"]:
         """Loads and returns a previously persisted ZinoState from a JSON file dump.
 

src/zino/utils.py

@@ -1,6 +1,9 @@
 import asyncio
+import logging
+from functools import wraps
 from ipaddress import ip_address
-from typing import Optional
+from time import time
+from typing import Optional, Union
 
 import aiodns
 from pyasn1.type.univ import OctetString
@@ -39,3 +42,24 @@ async def reverse_dns(ip: str) -> Optional[str]:
         return response.name
     except aiodns.error.DNSError:
         return None
+
+
+def log_time_spent(logger: Union[logging.Logger, str] = __name__, level: int = logging.DEBUG):
+    """Decorator that logs the time taken for a function to execute.  Not suitable for use with async functions"""
+    if isinstance(logger, str):
+        logger = logging.getLogger(logger)
+
+    def actual_decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start = time()
+            try:
+                result = func(*args, **kwargs)
+            finally:
+                end = time()
+                logger.log(level, "%s took %s seconds", func.__name__, end - start)
+            return result
+
+        return wrapper
+
+    return actual_decorator

tests/utils_test.py

@@ -1,10 +1,11 @@
+import logging
 from ipaddress import IPv4Address, IPv6Address
 from unittest.mock import AsyncMock, MagicMock
 
 import aiodns
 import pytest
 
-from zino.utils import parse_ip, reverse_dns
+from zino.utils import log_time_spent, parse_ip, reverse_dns
 
 
 class TestParseIP:
@@ -78,6 +79,24 @@ async def test_should_return_none_for_invalid_ip(self, mock_dnsresolver):
         assert result is None
 
 
+class TestLogTimeSpent:
+    def test_when_logger_is_specified_it_should_log_time_spent_using_that_logger(self, caplog):
+        @log_time_spent(logger="test_logger", level=logging.DEBUG)
+        def test_function():
+            pass
+
+        with caplog.at_level(logging.DEBUG):
+            test_function()
+
+        assert any(
+            record.name == "test_logger"
+            and "took" in record.msg
+            and "seconds" in record.msg
+            and "test_function" in record.args
+            for record in caplog.records
+        )
+
+
 @pytest.fixture
 def mock_dnsresolver(monkeypatch) -> AsyncMock:
     mock_dnsresolver = AsyncMock()