1+ # -*- encoding: utf-8 -*-
2+ import functools
3+ import threading
4+ from typing import Any , Callable , Dict , TypeVar , Type
5+
6+ F = TypeVar ('F' , bound = Callable [..., Any ])
7+
8+
9+ class ThreadSafeMeta (type ):
10+ """Metaclass that automatically adds thread safety to class methods."""
11+
12+ def __new__ (mcs , name : str , bases : tuple , namespace : Dict [str , Any ], ** kwargs ):
13+ # Create the class first
14+ cls = super ().__new__ (mcs , name , bases , namespace )
15+
16+ # Add a class-level lock if not already present
17+ if not hasattr (cls , '_lock' ):
18+ cls ._lock = threading .RLock ()
19+
20+ # Get methods that should be thread-safe (exclude private/dunder methods)
21+ thread_safe_methods = getattr (cls , '_thread_safe_methods' , None )
22+ if thread_safe_methods is None :
23+ # Auto-detect public methods that modify state
24+ thread_safe_methods = [
25+ method_name for method_name in namespace
26+ if (callable (getattr (cls , method_name , None )) and
27+ not method_name .startswith ('_' ) and
28+ method_name not in '__enter__' , '__exit__' , '__init__' ])
29+ ]
30+
31+ # Wrap each method with automatic locking
32+ for method_name in thread_safe_methods :
33+ if hasattr (cls , method_name ):
34+ original_method = getattr (cls , method_name )
35+ if callable (original_method ):
36+ wrapped_method = thread_safe (original_method )
37+ setattr (cls , method_name , wrapped_method )
38+
39+ return cls
40+
41+
42+ def thread_safe (func : F ) -> F :
43+ """Decorator that automatically adds thread safety to methods."""
44+
45+ @functools .wraps (func )
46+ def wrapper (self , * args , ** kwargs ):
47+ # Use instance lock if available, otherwise class lock
48+ lock = getattr (self , '_lock' , None )
49+ if lock is None :
50+ # Check if class has lock, if not create one
51+ if not hasattr (self .__class__ , '_lock' ):
52+ self .__class__ ._lock = threading .RLock ()
53+ lock = self .__class__ ._lock
54+
55+ with lock :
56+ return func (self , * args , ** kwargs )
57+
58+ return wrapper
59+
60+
61+ def auto_thread_safe (thread_safe_methods : list = None ):
62+ """Class decorator that adds automatic thread safety to specified methods."""
63+
64+ def decorator (cls : Type ) -> Type :
65+ # Add lock to class if not present
66+ if not hasattr (cls , '_lock' ):
67+ cls ._lock = threading .RLock ()
68+
69+ # Store thread-safe methods list
70+ if thread_safe_methods :
71+ cls ._thread_safe_methods = thread_safe_methods
72+
73+ # Get methods to make thread-safe
74+ methods_to_wrap = thread_safe_methods or [
75+ method_name for method_name in dir (cls )
76+ if (callable (getattr (cls , method_name , None )) and
77+ not method_name .startswith ('_' ) and
78+ method_name not in '__enter__' , '__exit__' , '__init__' ])
79+ ]
80+
81+ # Wrap each method
82+ for method_name in methods_to_wrap :
83+ if hasattr (cls , method_name ):
84+ original_method = getattr (cls , method_name )
85+ if callable (original_method ) and not hasattr (original_method , '_thread_safe_wrapped' ):
86+ wrapped_method = thread_safe (original_method )
87+ wrapped_method ._thread_safe_wrapped = True
88+ setattr (cls , method_name , wrapped_method )
89+
90+ return cls
91+
92+ return decorator
93+
94+
95+ class AutoThreadSafe :
96+ """Base class that provides automatic thread safety for all public methods."""
97+
98+ def __init__ (self ):
99+ if not hasattr (self , '_lock' ):
100+ self ._lock = threading .RLock ()
101+
102+ def __init_subclass__ (cls , ** kwargs ):
103+ super ().__init_subclass__ (** kwargs )
104+
105+ # Add class-level lock
106+ if not hasattr (cls , '_lock' ):
107+ cls ._lock = threading .RLock ()
108+
109+ # Auto-wrap public methods
110+ for attr_name in dir (cls ):
111+ if not attr_name .startswith ('_' ):
112+ attr = getattr (cls , attr_name )
113+ if callable (attr ) and not hasattr (attr , '_thread_safe_wrapped' ):
114+ wrapped_attr = thread_safe (attr )
115+ wrapped_attr ._thread_safe_wrapped = True
116+ setattr (cls , attr_name , wrapped_attr )
117+
118+
119+ def synchronized_method (func : F ) -> F :
120+ """Decorator for individual methods that need thread safety."""
121+ return thread_safe (func )
122+
123+
124+ class ThreadSafeContext :
125+ """Context manager for thread-safe operations."""
126+
127+ def __init__ (self , lock : threading .Lock ):
128+ self .lock = lock
129+
130+ def __enter__ (self ):
131+ self .lock .acquire ()
132+ return self
133+
134+ def __exit__ (self , exc_type , exc_val , exc_tb ):
135+ self .lock .release ()
0 commit comments