diegobeyl/backtesting

1"""
2Timeout decorators for Backtesting Application V2
3Provides @async_timeout and @sync_timeout for protecting I/O operations
4"""
5
6import asyncio
7import signal
8import functools
9from typing import TypeVar, Callable, Any, Optional
10import logging
11
12from utils.exceptions import TimeoutException
13
14logger = logging.getLogger(__name__)
15
16T = TypeVar('T')
17
18
19def async_timeout(seconds: float = 300) -> Callable:
20    """
21    Decorator to add timeout to async functions
22    
23    Args:
24        seconds: Timeout in seconds (default 300 = 5 minutes)
25        
26    Returns:
27        Decorated function that raises TimeoutException on timeout
28        
29    Example:
30        @async_timeout(30)
31        async def long_running_task():
32            await asyncio.sleep(60)  # Will timeout after 30s
33    """
34    def decorator(func: Callable) -> Callable:
35        @functools.wraps(func)
36        async def async_wrapper(*args, **kwargs) -> Any:
37            try:
38                return await asyncio.wait_for(
39                    func(*args, **kwargs),
40                    timeout=seconds
41                )
42            except asyncio.TimeoutError:
43                operation_name = getattr(func, '__name__', 'unknown_operation')
44                error_msg = f"Operation '{operation_name}' exceeded {seconds} second timeout"
45                logger.error(error_msg)
46                raise TimeoutException(
47                    error_msg,
48                    operation=operation_name,
49                    timeout_seconds=seconds
50                )
51            except TimeoutException:
52                raise
53            except Exception as e:
54                operation_name = getattr(func, '__name__', 'unknown_operation')
55                logger.error(f"Error in async_timeout wrapper for {operation_name}: {str(e)}")
56                raise
57        
58        return async_wrapper
59    
60    return decorator
61
62
63def sync_timeout(seconds: float = 30) -> Callable:
64    """
65    Decorator to add timeout to synchronous functions (UNIX/Linux only)
66    Note: Not supported on Windows. Falls back to no timeout on Windows.
67    
68    Args:
69        seconds: Timeout in seconds (default 30)
70        
71    Returns:
72        Decorated function that raises TimeoutException on timeout
73        
74    Example:
75        @sync_timeout(10)
76        def fetch_data():
77            time.sleep(20)  # Will timeout after 10s
78    """
79    def decorator(func: Callable) -> Callable:
80        @functools.wraps(func)
81        def sync_wrapper(*args, **kwargs) -> Any:
82            import platform
83            
84            # Check if signal.SIGALRM is available (Unix/Linux only)
85            if not hasattr(signal, 'SIGALRM'):
86                # Windows doesn't support signal.SIGALRM, so run without timeout
87                operation_name = getattr(func, '__name__', 'unknown_operation')
88                logger.warning(
89                    f"sync_timeout not supported on {platform.system()}, "
90                    f"running '{operation_name}' without timeout"
91                )
92                return func(*args, **kwargs)
93            
94            operation_name = getattr(func, '__name__', 'unknown_operation')
95            
96            def timeout_handler(signum, frame):
97                error_msg = f"Operation '{operation_name}' exceeded {seconds} second timeout"
98                raise TimeoutException(
99                    error_msg,
100                    operation=operation_name,
101                    timeout_seconds=seconds
102                )
103            
104            # Set up signal handler
105            old_handler = signal.signal(signal.SIGALRM, timeout_handler)
106            signal.alarm(int(seconds))
107            
108            try:
109                result = func(*args, **kwargs)
110                signal.alarm(0)  # Cancel the alarm
111                return result
112            except TimeoutException:
113                raise
114            except Exception as e:
115                signal.alarm(0)  # Cancel the alarm
116                logger.error(f"Error in sync_timeout wrapper for {operation_name}: {str(e)}")
117                raise
118            finally:
119                # Restore previous signal handler
120                signal.signal(signal.SIGALRM, old_handler)
121        
122        return sync_wrapper
123    
124    return decorator
125
126
127def with_timeout(timeout_seconds: Optional[float] = None) -> Callable:
128    """
129    Universal timeout decorator that works for both async and sync functions
130    Automatically detects the function type
131    
132    Args:
133        timeout_seconds: Timeout in seconds (None = no timeout)
134        
135    Returns:
136        Decorated function
137    """
138    def decorator(func: Callable) -> Callable:
139        if asyncio.iscoroutinefunction(func):
140            # It's an async function
141            timeout_val = timeout_seconds or 300
142            return async_timeout(timeout_val)(func)
143        else:
144            # It's a sync function
145            timeout_val = timeout_seconds or 30
146            return sync_timeout(timeout_val)(func)
147    
148    return decorator
149