stouputils 1.12.2__py3-none-any.whl → 1.13.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- stouputils/__main__.py +11 -6
- stouputils/continuous_delivery/pypi.py +39 -1
- stouputils/continuous_delivery/pypi.pyi +9 -0
- stouputils/ctx.py +408 -408
- stouputils/data_science/config/set.py +125 -125
- stouputils/data_science/models/keras_utils/callbacks/model_checkpoint_v2.py +31 -31
- stouputils/data_science/utils.py +285 -285
- stouputils/installer/__init__.py +18 -18
- stouputils/installer/linux.py +144 -144
- stouputils/installer/main.py +223 -223
- stouputils/installer/windows.py +136 -136
- stouputils/py.typed +1 -1
- stouputils/stouputils/__init__.pyi +15 -0
- stouputils/stouputils/_deprecated.pyi +12 -0
- stouputils/stouputils/all_doctests.pyi +46 -0
- stouputils/stouputils/applications/__init__.pyi +2 -0
- stouputils/stouputils/applications/automatic_docs.pyi +106 -0
- stouputils/stouputils/applications/upscaler/__init__.pyi +3 -0
- stouputils/stouputils/applications/upscaler/config.pyi +18 -0
- stouputils/stouputils/applications/upscaler/image.pyi +109 -0
- stouputils/stouputils/applications/upscaler/video.pyi +60 -0
- stouputils/stouputils/archive.pyi +67 -0
- stouputils/stouputils/backup.pyi +109 -0
- stouputils/stouputils/collections.pyi +86 -0
- stouputils/stouputils/continuous_delivery/__init__.pyi +5 -0
- stouputils/stouputils/continuous_delivery/cd_utils.pyi +129 -0
- stouputils/stouputils/continuous_delivery/github.pyi +162 -0
- stouputils/stouputils/continuous_delivery/pypi.pyi +53 -0
- stouputils/stouputils/continuous_delivery/pyproject.pyi +67 -0
- stouputils/stouputils/continuous_delivery/stubs.pyi +39 -0
- stouputils/stouputils/ctx.pyi +211 -0
- stouputils/stouputils/decorators.pyi +242 -0
- stouputils/stouputils/image.pyi +172 -0
- stouputils/stouputils/installer/__init__.pyi +5 -0
- stouputils/stouputils/installer/common.pyi +39 -0
- stouputils/stouputils/installer/downloader.pyi +24 -0
- stouputils/stouputils/installer/linux.pyi +39 -0
- stouputils/stouputils/installer/main.pyi +57 -0
- stouputils/stouputils/installer/windows.pyi +31 -0
- stouputils/stouputils/io.pyi +213 -0
- stouputils/stouputils/parallel.pyi +211 -0
- stouputils/stouputils/print.pyi +136 -0
- stouputils/stouputils/version_pkg.pyi +15 -0
- {stouputils-1.12.2.dist-info → stouputils-1.13.0.dist-info}/METADATA +1 -1
- {stouputils-1.12.2.dist-info → stouputils-1.13.0.dist-info}/RECORD +47 -16
- {stouputils-1.12.2.dist-info → stouputils-1.13.0.dist-info}/WHEEL +0 -0
- {stouputils-1.12.2.dist-info → stouputils-1.13.0.dist-info}/entry_points.txt +0 -0
stouputils/ctx.py
CHANGED
|
@@ -1,408 +1,408 @@
|
|
|
1
|
-
"""
|
|
2
|
-
This module provides context managers for various utilities such as logging to a file,
|
|
3
|
-
measuring execution time, silencing output, and setting multiprocessing start methods.
|
|
4
|
-
|
|
5
|
-
- LogToFile: Context manager to log to a file every print call (with LINE_UP handling)
|
|
6
|
-
- MeasureTime: Context manager to measure execution time of a code block
|
|
7
|
-
- Muffle: Context manager that temporarily silences output (alternative to stouputils.decorators.silent())
|
|
8
|
-
- DoNothing: Context manager that does nothing (no-op)
|
|
9
|
-
- SetMPStartMethod: Context manager to temporarily set multiprocessing start method
|
|
10
|
-
|
|
11
|
-
.. image:: https://raw.githubusercontent.com/Stoupy51/stouputils/refs/heads/main/assets/ctx_module.gif
|
|
12
|
-
:alt: stouputils ctx examples
|
|
13
|
-
"""
|
|
14
|
-
|
|
15
|
-
# Imports
|
|
16
|
-
from __future__ import annotations
|
|
17
|
-
|
|
18
|
-
import os
|
|
19
|
-
import sys
|
|
20
|
-
import time
|
|
21
|
-
from collections.abc import Callable
|
|
22
|
-
from contextlib import AbstractAsyncContextManager, AbstractContextManager
|
|
23
|
-
from typing import IO, Any, TextIO, TypeVar
|
|
24
|
-
|
|
25
|
-
from .io import super_open
|
|
26
|
-
from .print import TeeMultiOutput, debug
|
|
27
|
-
|
|
28
|
-
# Type variable for context managers
|
|
29
|
-
T = TypeVar("T")
|
|
30
|
-
|
|
31
|
-
# Abstract base class for context managers supporting both sync and async usage
|
|
32
|
-
class AbstractBothContextManager[T](AbstractContextManager[T], AbstractAsyncContextManager[T]):
|
|
33
|
-
""" Abstract base class for context managers that support both synchronous and asynchronous usage. """
|
|
34
|
-
pass
|
|
35
|
-
|
|
36
|
-
# Context manager to log to a file
|
|
37
|
-
class LogToFile(AbstractBothContextManager["LogToFile"]):
|
|
38
|
-
""" Context manager to log to a file.
|
|
39
|
-
|
|
40
|
-
This context manager allows you to temporarily log output to a file while still printing normally.
|
|
41
|
-
The file will receive log messages without ANSI color codes.
|
|
42
|
-
|
|
43
|
-
Args:
|
|
44
|
-
path (str): Path to the log file
|
|
45
|
-
mode (str): Mode to open the file in (default: "w")
|
|
46
|
-
encoding (str): Encoding to use for the file (default: "utf-8")
|
|
47
|
-
tee_stdout (bool): Whether to redirect stdout to the file (default: True)
|
|
48
|
-
tee_stderr (bool): Whether to redirect stderr to the file (default: True)
|
|
49
|
-
ignore_lineup (bool): Whether to ignore lines containing LINE_UP escape sequence in files (default: False)
|
|
50
|
-
restore_on_exit (bool): Whether to restore original stdout/stderr on exit (default: False)
|
|
51
|
-
This ctx uses TeeMultiOutput which handles closed files gracefully, so restoring is not mandatory.
|
|
52
|
-
|
|
53
|
-
Examples:
|
|
54
|
-
.. code-block:: python
|
|
55
|
-
|
|
56
|
-
> import stouputils as stp
|
|
57
|
-
> with stp.LogToFile("output.log"):
|
|
58
|
-
> stp.info("This will be logged to output.log and printed normally")
|
|
59
|
-
> print("This will also be logged")
|
|
60
|
-
|
|
61
|
-
> with stp.LogToFile("output.log") as log_ctx:
|
|
62
|
-
> stp.warning("This will be logged to output.log and printed normally")
|
|
63
|
-
> log_ctx.change_file("new_file.log")
|
|
64
|
-
> print("This will be logged to new_file.log")
|
|
65
|
-
"""
|
|
66
|
-
def __init__(
|
|
67
|
-
self,
|
|
68
|
-
path: str,
|
|
69
|
-
mode: str = "w",
|
|
70
|
-
encoding: str = "utf-8",
|
|
71
|
-
tee_stdout: bool = True,
|
|
72
|
-
tee_stderr: bool = True,
|
|
73
|
-
ignore_lineup: bool = True,
|
|
74
|
-
restore_on_exit: bool = False
|
|
75
|
-
) -> None:
|
|
76
|
-
self.path: str = path
|
|
77
|
-
""" Attribute remembering path to the log file """
|
|
78
|
-
self.mode: str = mode
|
|
79
|
-
""" Attribute remembering mode to open the file in """
|
|
80
|
-
self.encoding: str = encoding
|
|
81
|
-
""" Attribute remembering encoding to use for the file """
|
|
82
|
-
self.tee_stdout: bool = tee_stdout
|
|
83
|
-
""" Whether to redirect stdout to the file """
|
|
84
|
-
self.tee_stderr: bool = tee_stderr
|
|
85
|
-
""" Whether to redirect stderr to the file """
|
|
86
|
-
self.ignore_lineup: bool = ignore_lineup
|
|
87
|
-
""" Whether to ignore lines containing LINE_UP escape sequence in files """
|
|
88
|
-
self.restore_on_exit: bool = restore_on_exit
|
|
89
|
-
""" Whether to restore original stdout/stderr on exit.
|
|
90
|
-
This ctx uses TeeMultiOutput which handles closed files gracefully, so restoring is not mandatory. """
|
|
91
|
-
self.file: IO[Any]
|
|
92
|
-
""" Attribute remembering opened file """
|
|
93
|
-
self.original_stdout: TextIO
|
|
94
|
-
""" Original stdout before redirection """
|
|
95
|
-
self.original_stderr: TextIO
|
|
96
|
-
""" Original stderr before redirection """
|
|
97
|
-
|
|
98
|
-
def __enter__(self) -> LogToFile:
|
|
99
|
-
""" Enter context manager which opens the log file and redirects stdout/stderr """
|
|
100
|
-
# Open file
|
|
101
|
-
self.file = super_open(self.path, mode=self.mode, encoding=self.encoding)
|
|
102
|
-
|
|
103
|
-
# Redirect stdout and stderr if requested
|
|
104
|
-
if self.tee_stdout:
|
|
105
|
-
self.original_stdout = sys.stdout
|
|
106
|
-
sys.stdout = TeeMultiOutput(self.original_stdout, self.file, ignore_lineup=self.ignore_lineup)
|
|
107
|
-
if self.tee_stderr:
|
|
108
|
-
self.original_stderr = sys.stderr
|
|
109
|
-
sys.stderr = TeeMultiOutput(self.original_stderr, self.file, ignore_lineup=self.ignore_lineup)
|
|
110
|
-
|
|
111
|
-
# Return self
|
|
112
|
-
return self
|
|
113
|
-
|
|
114
|
-
def __exit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
115
|
-
""" Exit context manager which closes the log file and restores stdout/stderr """
|
|
116
|
-
# Restore original stdout and stderr (if requested)
|
|
117
|
-
if self.restore_on_exit:
|
|
118
|
-
if self.tee_stdout:
|
|
119
|
-
sys.stdout = self.original_stdout
|
|
120
|
-
if self.tee_stderr:
|
|
121
|
-
sys.stderr = self.original_stderr
|
|
122
|
-
|
|
123
|
-
# Close file
|
|
124
|
-
self.file.close()
|
|
125
|
-
|
|
126
|
-
async def __aenter__(self) -> LogToFile:
|
|
127
|
-
""" Enter async context manager which opens the log file and redirects stdout/stderr """
|
|
128
|
-
return self.__enter__()
|
|
129
|
-
|
|
130
|
-
async def __aexit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
131
|
-
""" Exit async context manager which closes the log file and restores stdout/stderr """
|
|
132
|
-
self.__exit__(exc_type, exc_val, exc_tb)
|
|
133
|
-
|
|
134
|
-
def change_file(self, new_path: str) -> None:
|
|
135
|
-
""" Change the log file to a new path.
|
|
136
|
-
|
|
137
|
-
Args:
|
|
138
|
-
new_path (str): New path to the log file
|
|
139
|
-
"""
|
|
140
|
-
# Close current file, open new file and redirect outputs
|
|
141
|
-
self.file.close()
|
|
142
|
-
self.path = new_path
|
|
143
|
-
self.__enter__()
|
|
144
|
-
|
|
145
|
-
@staticmethod
|
|
146
|
-
def common(logs_folder: str, filepath: str, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
|
|
147
|
-
""" Common code used at the beginning of a program to launch main function
|
|
148
|
-
|
|
149
|
-
Args:
|
|
150
|
-
logs_folder (str): Folder to store logs in
|
|
151
|
-
filepath (str): Path to the main function
|
|
152
|
-
func (Callable[..., Any]): Main function to launch
|
|
153
|
-
*args (tuple[Any, ...]): Arguments to pass to the main function
|
|
154
|
-
**kwargs (dict[str, Any]): Keyword arguments to pass to the main function
|
|
155
|
-
Returns:
|
|
156
|
-
Any: Return value of the main function
|
|
157
|
-
|
|
158
|
-
Examples:
|
|
159
|
-
>>> if __name__ == "__main__":
|
|
160
|
-
... LogToFile.common(f"{ROOT}/logs", __file__, main)
|
|
161
|
-
"""
|
|
162
|
-
# Import datetime
|
|
163
|
-
from datetime import datetime
|
|
164
|
-
|
|
165
|
-
# Build log file path
|
|
166
|
-
file_basename: str = os.path.splitext(os.path.basename(filepath))[0]
|
|
167
|
-
date_time: str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
|
168
|
-
date_str, time_str = date_time.split("_")
|
|
169
|
-
log_filepath: str = f"{logs_folder}/{file_basename}/{date_str}/{time_str}.log"
|
|
170
|
-
|
|
171
|
-
# Launch function with arguments if any
|
|
172
|
-
with LogToFile(log_filepath):
|
|
173
|
-
return func(*args, **kwargs)
|
|
174
|
-
|
|
175
|
-
# Context manager to measure execution time
|
|
176
|
-
class MeasureTime(AbstractBothContextManager["MeasureTime"]):
|
|
177
|
-
""" Context manager to measure execution time.
|
|
178
|
-
|
|
179
|
-
This context manager measures the execution time of the code block it wraps
|
|
180
|
-
and prints the result using a specified print function.
|
|
181
|
-
|
|
182
|
-
Args:
|
|
183
|
-
print_func (Callable): Function to use to print the execution time (e.g. debug, info, warning, error, etc.).
|
|
184
|
-
message (str): Message to display with the execution time. Defaults to "Execution time".
|
|
185
|
-
perf_counter (bool): Whether to use time.perf_counter_ns or time.time_ns. Defaults to True.
|
|
186
|
-
|
|
187
|
-
Examples:
|
|
188
|
-
.. code-block:: python
|
|
189
|
-
|
|
190
|
-
> import time
|
|
191
|
-
> import stouputils as stp
|
|
192
|
-
> with stp.MeasureTime(stp.info, message="My operation"):
|
|
193
|
-
... time.sleep(0.5)
|
|
194
|
-
> # [INFO HH:MM:SS] My operation: 500.123ms (500123456ns)
|
|
195
|
-
|
|
196
|
-
> with stp.MeasureTime(): # Uses debug by default
|
|
197
|
-
... time.sleep(0.1)
|
|
198
|
-
> # [DEBUG HH:MM:SS] Execution time: 100.456ms (100456789ns)
|
|
199
|
-
"""
|
|
200
|
-
def __init__(
|
|
201
|
-
self,
|
|
202
|
-
print_func: Callable[..., None] = debug,
|
|
203
|
-
message: str = "Execution time",
|
|
204
|
-
perf_counter: bool = True
|
|
205
|
-
) -> None:
|
|
206
|
-
self.print_func: Callable[..., None] = print_func
|
|
207
|
-
""" Function to use for printing the execution time """
|
|
208
|
-
self.message: str = message
|
|
209
|
-
""" Message to display with the execution time """
|
|
210
|
-
self.perf_counter: bool = perf_counter
|
|
211
|
-
""" Whether to use time.perf_counter_ns or time.time_ns """
|
|
212
|
-
self.ns: Callable[[], int] = time.perf_counter_ns if perf_counter else time.time_ns
|
|
213
|
-
""" Time function to use """
|
|
214
|
-
self.start_ns: int = 0
|
|
215
|
-
""" Start time in nanoseconds """
|
|
216
|
-
|
|
217
|
-
def __enter__(self) -> MeasureTime:
|
|
218
|
-
""" Enter context manager, record start time """
|
|
219
|
-
self.start_ns = self.ns()
|
|
220
|
-
return self
|
|
221
|
-
|
|
222
|
-
def __exit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
223
|
-
""" Exit context manager, calculate duration and print """
|
|
224
|
-
# Measure the execution time (nanoseconds and seconds)
|
|
225
|
-
total_ns: int = self.ns() - self.start_ns
|
|
226
|
-
total_ms: float = total_ns / 1_000_000
|
|
227
|
-
total_s: float = total_ns / 1_000_000_000
|
|
228
|
-
|
|
229
|
-
# Print the execution time (nanoseconds if less than 0.1s, seconds otherwise)
|
|
230
|
-
if total_ms < 100:
|
|
231
|
-
self.print_func(f"{self.message}: {total_ms:.3f}ms ({total_ns}ns)")
|
|
232
|
-
elif total_s < 60:
|
|
233
|
-
self.print_func(f"{self.message}: {(total_s):.5f}s")
|
|
234
|
-
else:
|
|
235
|
-
minutes: int = int(total_s) // 60
|
|
236
|
-
seconds: int = int(total_s) % 60
|
|
237
|
-
if minutes < 60:
|
|
238
|
-
self.print_func(f"{self.message}: {minutes}m {seconds}s")
|
|
239
|
-
else:
|
|
240
|
-
hours: int = minutes // 60
|
|
241
|
-
minutes: int = minutes % 60
|
|
242
|
-
if hours < 24:
|
|
243
|
-
self.print_func(f"{self.message}: {hours}h {minutes}m {seconds}s")
|
|
244
|
-
else:
|
|
245
|
-
days: int = hours // 24
|
|
246
|
-
hours: int = hours % 24
|
|
247
|
-
self.print_func(f"{self.message}: {days}d {hours}h {minutes}m {seconds}s")
|
|
248
|
-
|
|
249
|
-
async def __aenter__(self) -> MeasureTime:
|
|
250
|
-
""" Enter async context manager, record start time """
|
|
251
|
-
return self.__enter__()
|
|
252
|
-
|
|
253
|
-
async def __aexit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
254
|
-
""" Exit async context manager, calculate duration and print """
|
|
255
|
-
self.__exit__(exc_type, exc_val, exc_tb)
|
|
256
|
-
|
|
257
|
-
# Context manager to temporarily silence output
|
|
258
|
-
class Muffle(AbstractBothContextManager["Muffle"]):
|
|
259
|
-
""" Context manager that temporarily silences output.
|
|
260
|
-
(No thread-safety guaranteed)
|
|
261
|
-
|
|
262
|
-
Alternative to stouputils.decorators.silent()
|
|
263
|
-
|
|
264
|
-
Examples:
|
|
265
|
-
>>> with Muffle():
|
|
266
|
-
... print("This will not be printed")
|
|
267
|
-
"""
|
|
268
|
-
def __init__(self, mute_stderr: bool = False) -> None:
|
|
269
|
-
self.mute_stderr: bool = mute_stderr
|
|
270
|
-
""" Attribute remembering if stderr should be muted """
|
|
271
|
-
self.original_stdout: IO[Any]
|
|
272
|
-
""" Attribute remembering original stdout """
|
|
273
|
-
self.original_stderr: IO[Any]
|
|
274
|
-
""" Attribute remembering original stderr """
|
|
275
|
-
|
|
276
|
-
def __enter__(self) -> Muffle:
|
|
277
|
-
""" Enter context manager which redirects stdout and stderr to devnull """
|
|
278
|
-
# Redirect stdout to devnull
|
|
279
|
-
self.original_stdout = sys.stdout
|
|
280
|
-
sys.stdout = open(os.devnull, "w", encoding="utf-8")
|
|
281
|
-
|
|
282
|
-
# Redirect stderr to devnull if needed
|
|
283
|
-
if self.mute_stderr:
|
|
284
|
-
self.original_stderr = sys.stderr
|
|
285
|
-
sys.stderr = open(os.devnull, "w", encoding="utf-8")
|
|
286
|
-
|
|
287
|
-
# Return self
|
|
288
|
-
return self
|
|
289
|
-
|
|
290
|
-
def __exit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
291
|
-
""" Exit context manager which restores original stdout and stderr """
|
|
292
|
-
# Restore original stdout
|
|
293
|
-
sys.stdout.close()
|
|
294
|
-
sys.stdout = self.original_stdout
|
|
295
|
-
|
|
296
|
-
# Restore original stderr if needed
|
|
297
|
-
if self.mute_stderr:
|
|
298
|
-
sys.stderr.close()
|
|
299
|
-
sys.stderr = self.original_stderr
|
|
300
|
-
|
|
301
|
-
async def __aenter__(self) -> Muffle:
|
|
302
|
-
""" Enter async context manager which redirects stdout and stderr to devnull """
|
|
303
|
-
return self.__enter__()
|
|
304
|
-
|
|
305
|
-
async def __aexit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
306
|
-
""" Exit async context manager which restores original stdout and stderr """
|
|
307
|
-
self.__exit__(exc_type, exc_val, exc_tb)
|
|
308
|
-
|
|
309
|
-
# Context manager that does nothing
|
|
310
|
-
class DoNothing(AbstractBothContextManager["DoNothing"]):
|
|
311
|
-
""" Context manager that does nothing.
|
|
312
|
-
|
|
313
|
-
This is a no-op context manager that can be used as a placeholder
|
|
314
|
-
or for conditional context management.
|
|
315
|
-
|
|
316
|
-
Different from contextlib.nullcontext because it handles args and kwargs,
|
|
317
|
-
along with **async** context management.
|
|
318
|
-
|
|
319
|
-
Examples:
|
|
320
|
-
>>> with DoNothing():
|
|
321
|
-
... print("This will be printed normally")
|
|
322
|
-
This will be printed normally
|
|
323
|
-
|
|
324
|
-
>>> # Conditional context management
|
|
325
|
-
>>> some_condition = True
|
|
326
|
-
>>> ctx = DoNothing() if some_condition else Muffle()
|
|
327
|
-
>>> with ctx:
|
|
328
|
-
... print("May or may not be printed depending on condition")
|
|
329
|
-
May or may not be printed depending on condition
|
|
330
|
-
"""
|
|
331
|
-
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
332
|
-
""" No initialization needed, this is a no-op context manager """
|
|
333
|
-
pass
|
|
334
|
-
|
|
335
|
-
def __enter__(self) -> DoNothing:
|
|
336
|
-
""" Enter context manager (does nothing) """
|
|
337
|
-
return self
|
|
338
|
-
|
|
339
|
-
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
340
|
-
""" Exit context manager (does nothing) """
|
|
341
|
-
pass
|
|
342
|
-
|
|
343
|
-
async def __aenter__(self) -> DoNothing:
|
|
344
|
-
""" Enter async context manager (does nothing) """
|
|
345
|
-
return self
|
|
346
|
-
|
|
347
|
-
async def __aexit__(self, *excinfo: Any) -> None:
|
|
348
|
-
""" Exit async context manager (does nothing) """
|
|
349
|
-
pass
|
|
350
|
-
NullContextManager = DoNothing
|
|
351
|
-
""" Alias for DoNothing context manager """
|
|
352
|
-
|
|
353
|
-
# Context manager to temporarily set multiprocessing start method
|
|
354
|
-
class SetMPStartMethod(AbstractBothContextManager["SetMPStartMethod"]):
|
|
355
|
-
""" Context manager to temporarily set multiprocessing start method.
|
|
356
|
-
|
|
357
|
-
This context manager allows you to temporarily change the multiprocessing start method
|
|
358
|
-
and automatically restores the original method when exiting the context.
|
|
359
|
-
|
|
360
|
-
Args:
|
|
361
|
-
start_method (str): The start method to use: "spawn", "fork", or "forkserver"
|
|
362
|
-
|
|
363
|
-
Examples:
|
|
364
|
-
.. code-block:: python
|
|
365
|
-
|
|
366
|
-
> import multiprocessing as mp
|
|
367
|
-
> import stouputils as stp
|
|
368
|
-
> # Temporarily use spawn method
|
|
369
|
-
> with stp.SetMPStartMethod("spawn"):
|
|
370
|
-
> ... # Your multiprocessing code here
|
|
371
|
-
> ... pass
|
|
372
|
-
|
|
373
|
-
> # Original method is automatically restored
|
|
374
|
-
"""
|
|
375
|
-
def __init__(self, start_method: str | None) -> None:
|
|
376
|
-
self.start_method: str | None = start_method
|
|
377
|
-
""" The start method to use """
|
|
378
|
-
self.old_method: str | None = None
|
|
379
|
-
""" The original start method to restore """
|
|
380
|
-
|
|
381
|
-
def __enter__(self) -> SetMPStartMethod:
|
|
382
|
-
""" Enter context manager which sets the start method """
|
|
383
|
-
if self.start_method is None:
|
|
384
|
-
return self
|
|
385
|
-
import multiprocessing as mp
|
|
386
|
-
|
|
387
|
-
self.old_method = mp.get_start_method(allow_none=True)
|
|
388
|
-
if self.old_method != self.start_method:
|
|
389
|
-
mp.set_start_method(self.start_method, force=True)
|
|
390
|
-
return self
|
|
391
|
-
|
|
392
|
-
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
393
|
-
""" Exit context manager which restores the original start method """
|
|
394
|
-
if self.start_method is None:
|
|
395
|
-
return
|
|
396
|
-
import multiprocessing as mp
|
|
397
|
-
|
|
398
|
-
if self.old_method != self.start_method:
|
|
399
|
-
mp.set_start_method(self.old_method, force=True)
|
|
400
|
-
|
|
401
|
-
async def __aenter__(self) -> SetMPStartMethod:
|
|
402
|
-
""" Enter async context manager which sets the start method """
|
|
403
|
-
return self.__enter__()
|
|
404
|
-
|
|
405
|
-
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
406
|
-
""" Exit async context manager which restores the original start method """
|
|
407
|
-
self.__exit__(exc_type, exc_val, exc_tb)
|
|
408
|
-
|
|
1
|
+
"""
|
|
2
|
+
This module provides context managers for various utilities such as logging to a file,
|
|
3
|
+
measuring execution time, silencing output, and setting multiprocessing start methods.
|
|
4
|
+
|
|
5
|
+
- LogToFile: Context manager to log to a file every print call (with LINE_UP handling)
|
|
6
|
+
- MeasureTime: Context manager to measure execution time of a code block
|
|
7
|
+
- Muffle: Context manager that temporarily silences output (alternative to stouputils.decorators.silent())
|
|
8
|
+
- DoNothing: Context manager that does nothing (no-op)
|
|
9
|
+
- SetMPStartMethod: Context manager to temporarily set multiprocessing start method
|
|
10
|
+
|
|
11
|
+
.. image:: https://raw.githubusercontent.com/Stoupy51/stouputils/refs/heads/main/assets/ctx_module.gif
|
|
12
|
+
:alt: stouputils ctx examples
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
# Imports
|
|
16
|
+
from __future__ import annotations
|
|
17
|
+
|
|
18
|
+
import os
|
|
19
|
+
import sys
|
|
20
|
+
import time
|
|
21
|
+
from collections.abc import Callable
|
|
22
|
+
from contextlib import AbstractAsyncContextManager, AbstractContextManager
|
|
23
|
+
from typing import IO, Any, TextIO, TypeVar
|
|
24
|
+
|
|
25
|
+
from .io import super_open
|
|
26
|
+
from .print import TeeMultiOutput, debug
|
|
27
|
+
|
|
28
|
+
# Type variable for context managers
|
|
29
|
+
T = TypeVar("T")
|
|
30
|
+
|
|
31
|
+
# Abstract base class for context managers supporting both sync and async usage
|
|
32
|
+
class AbstractBothContextManager[T](AbstractContextManager[T], AbstractAsyncContextManager[T]):
|
|
33
|
+
""" Abstract base class for context managers that support both synchronous and asynchronous usage. """
|
|
34
|
+
pass
|
|
35
|
+
|
|
36
|
+
# Context manager to log to a file
|
|
37
|
+
class LogToFile(AbstractBothContextManager["LogToFile"]):
|
|
38
|
+
""" Context manager to log to a file.
|
|
39
|
+
|
|
40
|
+
This context manager allows you to temporarily log output to a file while still printing normally.
|
|
41
|
+
The file will receive log messages without ANSI color codes.
|
|
42
|
+
|
|
43
|
+
Args:
|
|
44
|
+
path (str): Path to the log file
|
|
45
|
+
mode (str): Mode to open the file in (default: "w")
|
|
46
|
+
encoding (str): Encoding to use for the file (default: "utf-8")
|
|
47
|
+
tee_stdout (bool): Whether to redirect stdout to the file (default: True)
|
|
48
|
+
tee_stderr (bool): Whether to redirect stderr to the file (default: True)
|
|
49
|
+
ignore_lineup (bool): Whether to ignore lines containing LINE_UP escape sequence in files (default: False)
|
|
50
|
+
restore_on_exit (bool): Whether to restore original stdout/stderr on exit (default: False)
|
|
51
|
+
This ctx uses TeeMultiOutput which handles closed files gracefully, so restoring is not mandatory.
|
|
52
|
+
|
|
53
|
+
Examples:
|
|
54
|
+
.. code-block:: python
|
|
55
|
+
|
|
56
|
+
> import stouputils as stp
|
|
57
|
+
> with stp.LogToFile("output.log"):
|
|
58
|
+
> stp.info("This will be logged to output.log and printed normally")
|
|
59
|
+
> print("This will also be logged")
|
|
60
|
+
|
|
61
|
+
> with stp.LogToFile("output.log") as log_ctx:
|
|
62
|
+
> stp.warning("This will be logged to output.log and printed normally")
|
|
63
|
+
> log_ctx.change_file("new_file.log")
|
|
64
|
+
> print("This will be logged to new_file.log")
|
|
65
|
+
"""
|
|
66
|
+
def __init__(
|
|
67
|
+
self,
|
|
68
|
+
path: str,
|
|
69
|
+
mode: str = "w",
|
|
70
|
+
encoding: str = "utf-8",
|
|
71
|
+
tee_stdout: bool = True,
|
|
72
|
+
tee_stderr: bool = True,
|
|
73
|
+
ignore_lineup: bool = True,
|
|
74
|
+
restore_on_exit: bool = False
|
|
75
|
+
) -> None:
|
|
76
|
+
self.path: str = path
|
|
77
|
+
""" Attribute remembering path to the log file """
|
|
78
|
+
self.mode: str = mode
|
|
79
|
+
""" Attribute remembering mode to open the file in """
|
|
80
|
+
self.encoding: str = encoding
|
|
81
|
+
""" Attribute remembering encoding to use for the file """
|
|
82
|
+
self.tee_stdout: bool = tee_stdout
|
|
83
|
+
""" Whether to redirect stdout to the file """
|
|
84
|
+
self.tee_stderr: bool = tee_stderr
|
|
85
|
+
""" Whether to redirect stderr to the file """
|
|
86
|
+
self.ignore_lineup: bool = ignore_lineup
|
|
87
|
+
""" Whether to ignore lines containing LINE_UP escape sequence in files """
|
|
88
|
+
self.restore_on_exit: bool = restore_on_exit
|
|
89
|
+
""" Whether to restore original stdout/stderr on exit.
|
|
90
|
+
This ctx uses TeeMultiOutput which handles closed files gracefully, so restoring is not mandatory. """
|
|
91
|
+
self.file: IO[Any]
|
|
92
|
+
""" Attribute remembering opened file """
|
|
93
|
+
self.original_stdout: TextIO
|
|
94
|
+
""" Original stdout before redirection """
|
|
95
|
+
self.original_stderr: TextIO
|
|
96
|
+
""" Original stderr before redirection """
|
|
97
|
+
|
|
98
|
+
def __enter__(self) -> LogToFile:
|
|
99
|
+
""" Enter context manager which opens the log file and redirects stdout/stderr """
|
|
100
|
+
# Open file
|
|
101
|
+
self.file = super_open(self.path, mode=self.mode, encoding=self.encoding)
|
|
102
|
+
|
|
103
|
+
# Redirect stdout and stderr if requested
|
|
104
|
+
if self.tee_stdout:
|
|
105
|
+
self.original_stdout = sys.stdout
|
|
106
|
+
sys.stdout = TeeMultiOutput(self.original_stdout, self.file, ignore_lineup=self.ignore_lineup)
|
|
107
|
+
if self.tee_stderr:
|
|
108
|
+
self.original_stderr = sys.stderr
|
|
109
|
+
sys.stderr = TeeMultiOutput(self.original_stderr, self.file, ignore_lineup=self.ignore_lineup)
|
|
110
|
+
|
|
111
|
+
# Return self
|
|
112
|
+
return self
|
|
113
|
+
|
|
114
|
+
def __exit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
115
|
+
""" Exit context manager which closes the log file and restores stdout/stderr """
|
|
116
|
+
# Restore original stdout and stderr (if requested)
|
|
117
|
+
if self.restore_on_exit:
|
|
118
|
+
if self.tee_stdout:
|
|
119
|
+
sys.stdout = self.original_stdout
|
|
120
|
+
if self.tee_stderr:
|
|
121
|
+
sys.stderr = self.original_stderr
|
|
122
|
+
|
|
123
|
+
# Close file
|
|
124
|
+
self.file.close()
|
|
125
|
+
|
|
126
|
+
async def __aenter__(self) -> LogToFile:
|
|
127
|
+
""" Enter async context manager which opens the log file and redirects stdout/stderr """
|
|
128
|
+
return self.__enter__()
|
|
129
|
+
|
|
130
|
+
async def __aexit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
131
|
+
""" Exit async context manager which closes the log file and restores stdout/stderr """
|
|
132
|
+
self.__exit__(exc_type, exc_val, exc_tb)
|
|
133
|
+
|
|
134
|
+
def change_file(self, new_path: str) -> None:
|
|
135
|
+
""" Change the log file to a new path.
|
|
136
|
+
|
|
137
|
+
Args:
|
|
138
|
+
new_path (str): New path to the log file
|
|
139
|
+
"""
|
|
140
|
+
# Close current file, open new file and redirect outputs
|
|
141
|
+
self.file.close()
|
|
142
|
+
self.path = new_path
|
|
143
|
+
self.__enter__()
|
|
144
|
+
|
|
145
|
+
@staticmethod
|
|
146
|
+
def common(logs_folder: str, filepath: str, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
|
|
147
|
+
""" Common code used at the beginning of a program to launch main function
|
|
148
|
+
|
|
149
|
+
Args:
|
|
150
|
+
logs_folder (str): Folder to store logs in
|
|
151
|
+
filepath (str): Path to the main function
|
|
152
|
+
func (Callable[..., Any]): Main function to launch
|
|
153
|
+
*args (tuple[Any, ...]): Arguments to pass to the main function
|
|
154
|
+
**kwargs (dict[str, Any]): Keyword arguments to pass to the main function
|
|
155
|
+
Returns:
|
|
156
|
+
Any: Return value of the main function
|
|
157
|
+
|
|
158
|
+
Examples:
|
|
159
|
+
>>> if __name__ == "__main__":
|
|
160
|
+
... LogToFile.common(f"{ROOT}/logs", __file__, main)
|
|
161
|
+
"""
|
|
162
|
+
# Import datetime
|
|
163
|
+
from datetime import datetime
|
|
164
|
+
|
|
165
|
+
# Build log file path
|
|
166
|
+
file_basename: str = os.path.splitext(os.path.basename(filepath))[0]
|
|
167
|
+
date_time: str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
|
168
|
+
date_str, time_str = date_time.split("_")
|
|
169
|
+
log_filepath: str = f"{logs_folder}/{file_basename}/{date_str}/{time_str}.log"
|
|
170
|
+
|
|
171
|
+
# Launch function with arguments if any
|
|
172
|
+
with LogToFile(log_filepath):
|
|
173
|
+
return func(*args, **kwargs)
|
|
174
|
+
|
|
175
|
+
# Context manager to measure execution time
|
|
176
|
+
class MeasureTime(AbstractBothContextManager["MeasureTime"]):
|
|
177
|
+
""" Context manager to measure execution time.
|
|
178
|
+
|
|
179
|
+
This context manager measures the execution time of the code block it wraps
|
|
180
|
+
and prints the result using a specified print function.
|
|
181
|
+
|
|
182
|
+
Args:
|
|
183
|
+
print_func (Callable): Function to use to print the execution time (e.g. debug, info, warning, error, etc.).
|
|
184
|
+
message (str): Message to display with the execution time. Defaults to "Execution time".
|
|
185
|
+
perf_counter (bool): Whether to use time.perf_counter_ns or time.time_ns. Defaults to True.
|
|
186
|
+
|
|
187
|
+
Examples:
|
|
188
|
+
.. code-block:: python
|
|
189
|
+
|
|
190
|
+
> import time
|
|
191
|
+
> import stouputils as stp
|
|
192
|
+
> with stp.MeasureTime(stp.info, message="My operation"):
|
|
193
|
+
... time.sleep(0.5)
|
|
194
|
+
> # [INFO HH:MM:SS] My operation: 500.123ms (500123456ns)
|
|
195
|
+
|
|
196
|
+
> with stp.MeasureTime(): # Uses debug by default
|
|
197
|
+
... time.sleep(0.1)
|
|
198
|
+
> # [DEBUG HH:MM:SS] Execution time: 100.456ms (100456789ns)
|
|
199
|
+
"""
|
|
200
|
+
def __init__(
|
|
201
|
+
self,
|
|
202
|
+
print_func: Callable[..., None] = debug,
|
|
203
|
+
message: str = "Execution time",
|
|
204
|
+
perf_counter: bool = True
|
|
205
|
+
) -> None:
|
|
206
|
+
self.print_func: Callable[..., None] = print_func
|
|
207
|
+
""" Function to use for printing the execution time """
|
|
208
|
+
self.message: str = message
|
|
209
|
+
""" Message to display with the execution time """
|
|
210
|
+
self.perf_counter: bool = perf_counter
|
|
211
|
+
""" Whether to use time.perf_counter_ns or time.time_ns """
|
|
212
|
+
self.ns: Callable[[], int] = time.perf_counter_ns if perf_counter else time.time_ns
|
|
213
|
+
""" Time function to use """
|
|
214
|
+
self.start_ns: int = 0
|
|
215
|
+
""" Start time in nanoseconds """
|
|
216
|
+
|
|
217
|
+
def __enter__(self) -> MeasureTime:
|
|
218
|
+
""" Enter context manager, record start time """
|
|
219
|
+
self.start_ns = self.ns()
|
|
220
|
+
return self
|
|
221
|
+
|
|
222
|
+
def __exit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
223
|
+
""" Exit context manager, calculate duration and print """
|
|
224
|
+
# Measure the execution time (nanoseconds and seconds)
|
|
225
|
+
total_ns: int = self.ns() - self.start_ns
|
|
226
|
+
total_ms: float = total_ns / 1_000_000
|
|
227
|
+
total_s: float = total_ns / 1_000_000_000
|
|
228
|
+
|
|
229
|
+
# Print the execution time (nanoseconds if less than 0.1s, seconds otherwise)
|
|
230
|
+
if total_ms < 100:
|
|
231
|
+
self.print_func(f"{self.message}: {total_ms:.3f}ms ({total_ns}ns)")
|
|
232
|
+
elif total_s < 60:
|
|
233
|
+
self.print_func(f"{self.message}: {(total_s):.5f}s")
|
|
234
|
+
else:
|
|
235
|
+
minutes: int = int(total_s) // 60
|
|
236
|
+
seconds: int = int(total_s) % 60
|
|
237
|
+
if minutes < 60:
|
|
238
|
+
self.print_func(f"{self.message}: {minutes}m {seconds}s")
|
|
239
|
+
else:
|
|
240
|
+
hours: int = minutes // 60
|
|
241
|
+
minutes: int = minutes % 60
|
|
242
|
+
if hours < 24:
|
|
243
|
+
self.print_func(f"{self.message}: {hours}h {minutes}m {seconds}s")
|
|
244
|
+
else:
|
|
245
|
+
days: int = hours // 24
|
|
246
|
+
hours: int = hours % 24
|
|
247
|
+
self.print_func(f"{self.message}: {days}d {hours}h {minutes}m {seconds}s")
|
|
248
|
+
|
|
249
|
+
async def __aenter__(self) -> MeasureTime:
|
|
250
|
+
""" Enter async context manager, record start time """
|
|
251
|
+
return self.__enter__()
|
|
252
|
+
|
|
253
|
+
async def __aexit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
254
|
+
""" Exit async context manager, calculate duration and print """
|
|
255
|
+
self.__exit__(exc_type, exc_val, exc_tb)
|
|
256
|
+
|
|
257
|
+
# Context manager to temporarily silence output
|
|
258
|
+
class Muffle(AbstractBothContextManager["Muffle"]):
|
|
259
|
+
""" Context manager that temporarily silences output.
|
|
260
|
+
(No thread-safety guaranteed)
|
|
261
|
+
|
|
262
|
+
Alternative to stouputils.decorators.silent()
|
|
263
|
+
|
|
264
|
+
Examples:
|
|
265
|
+
>>> with Muffle():
|
|
266
|
+
... print("This will not be printed")
|
|
267
|
+
"""
|
|
268
|
+
def __init__(self, mute_stderr: bool = False) -> None:
|
|
269
|
+
self.mute_stderr: bool = mute_stderr
|
|
270
|
+
""" Attribute remembering if stderr should be muted """
|
|
271
|
+
self.original_stdout: IO[Any]
|
|
272
|
+
""" Attribute remembering original stdout """
|
|
273
|
+
self.original_stderr: IO[Any]
|
|
274
|
+
""" Attribute remembering original stderr """
|
|
275
|
+
|
|
276
|
+
def __enter__(self) -> Muffle:
|
|
277
|
+
""" Enter context manager which redirects stdout and stderr to devnull """
|
|
278
|
+
# Redirect stdout to devnull
|
|
279
|
+
self.original_stdout = sys.stdout
|
|
280
|
+
sys.stdout = open(os.devnull, "w", encoding="utf-8")
|
|
281
|
+
|
|
282
|
+
# Redirect stderr to devnull if needed
|
|
283
|
+
if self.mute_stderr:
|
|
284
|
+
self.original_stderr = sys.stderr
|
|
285
|
+
sys.stderr = open(os.devnull, "w", encoding="utf-8")
|
|
286
|
+
|
|
287
|
+
# Return self
|
|
288
|
+
return self
|
|
289
|
+
|
|
290
|
+
def __exit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
291
|
+
""" Exit context manager which restores original stdout and stderr """
|
|
292
|
+
# Restore original stdout
|
|
293
|
+
sys.stdout.close()
|
|
294
|
+
sys.stdout = self.original_stdout
|
|
295
|
+
|
|
296
|
+
# Restore original stderr if needed
|
|
297
|
+
if self.mute_stderr:
|
|
298
|
+
sys.stderr.close()
|
|
299
|
+
sys.stderr = self.original_stderr
|
|
300
|
+
|
|
301
|
+
async def __aenter__(self) -> Muffle:
|
|
302
|
+
""" Enter async context manager which redirects stdout and stderr to devnull """
|
|
303
|
+
return self.__enter__()
|
|
304
|
+
|
|
305
|
+
async def __aexit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
|
|
306
|
+
""" Exit async context manager which restores original stdout and stderr """
|
|
307
|
+
self.__exit__(exc_type, exc_val, exc_tb)
|
|
308
|
+
|
|
309
|
+
# Context manager that does nothing
|
|
310
|
+
class DoNothing(AbstractBothContextManager["DoNothing"]):
|
|
311
|
+
""" Context manager that does nothing.
|
|
312
|
+
|
|
313
|
+
This is a no-op context manager that can be used as a placeholder
|
|
314
|
+
or for conditional context management.
|
|
315
|
+
|
|
316
|
+
Different from contextlib.nullcontext because it handles args and kwargs,
|
|
317
|
+
along with **async** context management.
|
|
318
|
+
|
|
319
|
+
Examples:
|
|
320
|
+
>>> with DoNothing():
|
|
321
|
+
... print("This will be printed normally")
|
|
322
|
+
This will be printed normally
|
|
323
|
+
|
|
324
|
+
>>> # Conditional context management
|
|
325
|
+
>>> some_condition = True
|
|
326
|
+
>>> ctx = DoNothing() if some_condition else Muffle()
|
|
327
|
+
>>> with ctx:
|
|
328
|
+
... print("May or may not be printed depending on condition")
|
|
329
|
+
May or may not be printed depending on condition
|
|
330
|
+
"""
|
|
331
|
+
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
332
|
+
""" No initialization needed, this is a no-op context manager """
|
|
333
|
+
pass
|
|
334
|
+
|
|
335
|
+
def __enter__(self) -> DoNothing:
|
|
336
|
+
""" Enter context manager (does nothing) """
|
|
337
|
+
return self
|
|
338
|
+
|
|
339
|
+
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
340
|
+
""" Exit context manager (does nothing) """
|
|
341
|
+
pass
|
|
342
|
+
|
|
343
|
+
async def __aenter__(self) -> DoNothing:
|
|
344
|
+
""" Enter async context manager (does nothing) """
|
|
345
|
+
return self
|
|
346
|
+
|
|
347
|
+
async def __aexit__(self, *excinfo: Any) -> None:
|
|
348
|
+
""" Exit async context manager (does nothing) """
|
|
349
|
+
pass
|
|
350
|
+
NullContextManager = DoNothing
|
|
351
|
+
""" Alias for DoNothing context manager """
|
|
352
|
+
|
|
353
|
+
# Context manager to temporarily set multiprocessing start method
|
|
354
|
+
class SetMPStartMethod(AbstractBothContextManager["SetMPStartMethod"]):
|
|
355
|
+
""" Context manager to temporarily set multiprocessing start method.
|
|
356
|
+
|
|
357
|
+
This context manager allows you to temporarily change the multiprocessing start method
|
|
358
|
+
and automatically restores the original method when exiting the context.
|
|
359
|
+
|
|
360
|
+
Args:
|
|
361
|
+
start_method (str): The start method to use: "spawn", "fork", or "forkserver"
|
|
362
|
+
|
|
363
|
+
Examples:
|
|
364
|
+
.. code-block:: python
|
|
365
|
+
|
|
366
|
+
> import multiprocessing as mp
|
|
367
|
+
> import stouputils as stp
|
|
368
|
+
> # Temporarily use spawn method
|
|
369
|
+
> with stp.SetMPStartMethod("spawn"):
|
|
370
|
+
> ... # Your multiprocessing code here
|
|
371
|
+
> ... pass
|
|
372
|
+
|
|
373
|
+
> # Original method is automatically restored
|
|
374
|
+
"""
|
|
375
|
+
def __init__(self, start_method: str | None) -> None:
|
|
376
|
+
self.start_method: str | None = start_method
|
|
377
|
+
""" The start method to use """
|
|
378
|
+
self.old_method: str | None = None
|
|
379
|
+
""" The original start method to restore """
|
|
380
|
+
|
|
381
|
+
def __enter__(self) -> SetMPStartMethod:
|
|
382
|
+
""" Enter context manager which sets the start method """
|
|
383
|
+
if self.start_method is None:
|
|
384
|
+
return self
|
|
385
|
+
import multiprocessing as mp
|
|
386
|
+
|
|
387
|
+
self.old_method = mp.get_start_method(allow_none=True)
|
|
388
|
+
if self.old_method != self.start_method:
|
|
389
|
+
mp.set_start_method(self.start_method, force=True)
|
|
390
|
+
return self
|
|
391
|
+
|
|
392
|
+
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
393
|
+
""" Exit context manager which restores the original start method """
|
|
394
|
+
if self.start_method is None:
|
|
395
|
+
return
|
|
396
|
+
import multiprocessing as mp
|
|
397
|
+
|
|
398
|
+
if self.old_method != self.start_method:
|
|
399
|
+
mp.set_start_method(self.old_method, force=True)
|
|
400
|
+
|
|
401
|
+
async def __aenter__(self) -> SetMPStartMethod:
|
|
402
|
+
""" Enter async context manager which sets the start method """
|
|
403
|
+
return self.__enter__()
|
|
404
|
+
|
|
405
|
+
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
406
|
+
""" Exit async context manager which restores the original start method """
|
|
407
|
+
self.__exit__(exc_type, exc_val, exc_tb)
|
|
408
|
+
|