From 0dc3394ba46e6f3a0d6438f7cbb6f361e1468021 Mon Sep 17 00:00:00 2001 From: Florian Rau Date: Thu, 5 Oct 2023 11:46:00 +0200 Subject: [PATCH 1/4] add @sleepless decorator --- ibllib/pipes/misc.py | 49 ++++++++++++++++++++++++++++++++++++-------- 1 file changed, 41 insertions(+), 8 deletions(-) diff --git a/ibllib/pipes/misc.py b/ibllib/pipes/misc.py index d3911533f..893b939b0 100644 --- a/ibllib/pipes/misc.py +++ b/ibllib/pipes/misc.py @@ -9,7 +9,7 @@ import time import logging from pathlib import Path -from typing import Union, List +from typing import Union, List, Callable, Any from inspect import signature import uuid import socket @@ -1147,13 +1147,46 @@ class WindowsInhibitor: ES_CONTINUOUS = 0x80000000 ES_SYSTEM_REQUIRED = 0x00000001 - def __init__(self): - pass - - def inhibit(self): - print("Preventing Windows from going to sleep") + @staticmethod + def inhibit(quiet: bool = False): + if quiet: + log.debug("Preventing Windows from going to sleep") + else: + print("Preventing Windows from going to sleep") ctypes.windll.kernel32.SetThreadExecutionState(WindowsInhibitor.ES_CONTINUOUS | WindowsInhibitor.ES_SYSTEM_REQUIRED) - def uninhibit(self): - print("Allowing Windows to go to sleep") + @staticmethod + def uninhibit(quiet: bool = False): + if quiet: + log.debug("Allowing Windows from going to sleep") + else: + print("Allowing Windows to go to sleep") ctypes.windll.kernel32.SetThreadExecutionState(WindowsInhibitor.ES_CONTINUOUS) + + +def sleepless(func: Callable[..., Any]) -> Callable[..., Any]: + """ + Decorator to ensure that the system doesn't enter sleep or idle mode during a long-running task. + + This decorator wraps a function and sets the thread execution state to prevent + the system from entering sleep or idle mode while the decorated function is + running. + + Parameters + ---------- + func : callable + The function to decorate. + + Returns + ------- + callable + The decorated function. + """ + def inner(*args, **kwargs) -> Any: + if os.name == 'nt': + WindowsInhibitor().inhibit(quiet=True) + result = func(*args, **kwargs) + if os.name == 'nt': + WindowsInhibitor().uninhibit(quiet=True) + return result + return inner From 92632dcd8cc574dffba6d38aa135d5bf52bf80c5 Mon Sep 17 00:00:00 2001 From: Florian Rau Date: Thu, 5 Oct 2023 11:56:19 +0200 Subject: [PATCH 2/4] refactor / add error log for failure --- ibllib/pipes/misc.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/ibllib/pipes/misc.py b/ibllib/pipes/misc.py index 893b939b0..17e694110 100644 --- a/ibllib/pipes/misc.py +++ b/ibllib/pipes/misc.py @@ -1147,13 +1147,19 @@ class WindowsInhibitor: ES_CONTINUOUS = 0x80000000 ES_SYSTEM_REQUIRED = 0x00000001 + @staticmethod + def _set_thread_execution_state(state: int) -> None: + result = ctypes.windll.kernel32.SetThreadExecutionState(state) + if result == 0: + log.error("Failed to set thread execution state.") + @staticmethod def inhibit(quiet: bool = False): if quiet: log.debug("Preventing Windows from going to sleep") else: print("Preventing Windows from going to sleep") - ctypes.windll.kernel32.SetThreadExecutionState(WindowsInhibitor.ES_CONTINUOUS | WindowsInhibitor.ES_SYSTEM_REQUIRED) + WindowsInhibitor._set_thread_execution_state(WindowsInhibitor.ES_CONTINUOUS | WindowsInhibitor.ES_SYSTEM_REQUIRED) @staticmethod def uninhibit(quiet: bool = False): @@ -1161,7 +1167,7 @@ def uninhibit(quiet: bool = False): log.debug("Allowing Windows from going to sleep") else: print("Allowing Windows to go to sleep") - ctypes.windll.kernel32.SetThreadExecutionState(WindowsInhibitor.ES_CONTINUOUS) + WindowsInhibitor._set_thread_execution_state(WindowsInhibitor.ES_CONTINUOUS) def sleepless(func: Callable[..., Any]) -> Callable[..., Any]: From c04749e179dba9c225be83f045f2eab1b2ef1d88 Mon Sep 17 00:00:00 2001 From: Florian Rau Date: Mon, 13 Nov 2023 16:38:39 +0000 Subject: [PATCH 3/4] Update misc.py --- ibllib/pipes/misc.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ibllib/pipes/misc.py b/ibllib/pipes/misc.py index 17e694110..d652f2ffc 100644 --- a/ibllib/pipes/misc.py +++ b/ibllib/pipes/misc.py @@ -8,6 +8,7 @@ import sys import time import logging +from functools import wraps from pathlib import Path from typing import Union, List, Callable, Any from inspect import signature @@ -1164,7 +1165,7 @@ def inhibit(quiet: bool = False): @staticmethod def uninhibit(quiet: bool = False): if quiet: - log.debug("Allowing Windows from going to sleep") + log.debug("Allowing Windows to go to sleep") else: print("Allowing Windows to go to sleep") WindowsInhibitor._set_thread_execution_state(WindowsInhibitor.ES_CONTINUOUS) @@ -1188,6 +1189,8 @@ def sleepless(func: Callable[..., Any]) -> Callable[..., Any]: callable The decorated function. """ + + @wraps(func) def inner(*args, **kwargs) -> Any: if os.name == 'nt': WindowsInhibitor().inhibit(quiet=True) From c19a9c750cc3f0a66593ad76d5391a891e8f086c Mon Sep 17 00:00:00 2001 From: Florian Rau Date: Mon, 13 Nov 2023 16:49:42 +0000 Subject: [PATCH 4/4] add test for sleepless decorator --- ibllib/tests/test_pipes.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/ibllib/tests/test_pipes.py b/ibllib/tests/test_pipes.py index ba5c282dd..cbe86462a 100644 --- a/ibllib/tests/test_pipes.py +++ b/ibllib/tests/test_pipes.py @@ -21,6 +21,7 @@ import ibllib.io.extractors.base import ibllib.tests.fixtures.utils as fu from ibllib.pipes import misc +from ibllib.pipes.misc import sleepless from ibllib.tests import TEST_DB import ibllib.pipes.scan_fix_passive_files as fix from ibllib.pipes.base_tasks import RegisterRawDataTask @@ -698,5 +699,23 @@ def test_rename_files(self): self.assertCountEqual(expected, files) +class TestSleeplessDecorator(unittest.TestCase): + + def test_decorator_argument_passing(self): + + def dummy_function(arg1, arg2): + return arg1, arg2 + + # Applying the decorator to the dummy function + decorated_func = sleepless(dummy_function) + + # Check if the function name is maintained + self.assertEqual(decorated_func.__name__, 'dummy_function') + + # Check if arguments are passed correctly + result = decorated_func("test1", "test2") + self.assertEqual(result, ("test1", "test2")) + + if __name__ == '__main__': unittest.main(exit=False, verbosity=2)