stouputils 1.12.1__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.
Files changed (138) hide show
  1. stouputils/__init__.py +40 -0
  2. stouputils/__init__.pyi +14 -0
  3. stouputils/__main__.py +81 -0
  4. stouputils/_deprecated.py +37 -0
  5. stouputils/_deprecated.pyi +12 -0
  6. stouputils/all_doctests.py +160 -0
  7. stouputils/all_doctests.pyi +46 -0
  8. stouputils/applications/__init__.py +22 -0
  9. stouputils/applications/__init__.pyi +2 -0
  10. stouputils/applications/automatic_docs.py +634 -0
  11. stouputils/applications/automatic_docs.pyi +106 -0
  12. stouputils/applications/upscaler/__init__.py +39 -0
  13. stouputils/applications/upscaler/__init__.pyi +3 -0
  14. stouputils/applications/upscaler/config.py +128 -0
  15. stouputils/applications/upscaler/config.pyi +18 -0
  16. stouputils/applications/upscaler/image.py +247 -0
  17. stouputils/applications/upscaler/image.pyi +109 -0
  18. stouputils/applications/upscaler/video.py +287 -0
  19. stouputils/applications/upscaler/video.pyi +60 -0
  20. stouputils/archive.py +344 -0
  21. stouputils/archive.pyi +67 -0
  22. stouputils/backup.py +488 -0
  23. stouputils/backup.pyi +109 -0
  24. stouputils/collections.py +244 -0
  25. stouputils/collections.pyi +86 -0
  26. stouputils/continuous_delivery/__init__.py +27 -0
  27. stouputils/continuous_delivery/__init__.pyi +5 -0
  28. stouputils/continuous_delivery/cd_utils.py +243 -0
  29. stouputils/continuous_delivery/cd_utils.pyi +129 -0
  30. stouputils/continuous_delivery/github.py +522 -0
  31. stouputils/continuous_delivery/github.pyi +162 -0
  32. stouputils/continuous_delivery/pypi.py +91 -0
  33. stouputils/continuous_delivery/pypi.pyi +43 -0
  34. stouputils/continuous_delivery/pyproject.py +147 -0
  35. stouputils/continuous_delivery/pyproject.pyi +67 -0
  36. stouputils/continuous_delivery/stubs.py +86 -0
  37. stouputils/continuous_delivery/stubs.pyi +39 -0
  38. stouputils/ctx.py +408 -0
  39. stouputils/ctx.pyi +211 -0
  40. stouputils/data_science/config/get.py +51 -0
  41. stouputils/data_science/config/set.py +125 -0
  42. stouputils/data_science/data_processing/image/__init__.py +66 -0
  43. stouputils/data_science/data_processing/image/auto_contrast.py +79 -0
  44. stouputils/data_science/data_processing/image/axis_flip.py +58 -0
  45. stouputils/data_science/data_processing/image/bias_field_correction.py +74 -0
  46. stouputils/data_science/data_processing/image/binary_threshold.py +73 -0
  47. stouputils/data_science/data_processing/image/blur.py +59 -0
  48. stouputils/data_science/data_processing/image/brightness.py +54 -0
  49. stouputils/data_science/data_processing/image/canny.py +110 -0
  50. stouputils/data_science/data_processing/image/clahe.py +92 -0
  51. stouputils/data_science/data_processing/image/common.py +30 -0
  52. stouputils/data_science/data_processing/image/contrast.py +53 -0
  53. stouputils/data_science/data_processing/image/curvature_flow_filter.py +74 -0
  54. stouputils/data_science/data_processing/image/denoise.py +378 -0
  55. stouputils/data_science/data_processing/image/histogram_equalization.py +123 -0
  56. stouputils/data_science/data_processing/image/invert.py +64 -0
  57. stouputils/data_science/data_processing/image/laplacian.py +60 -0
  58. stouputils/data_science/data_processing/image/median_blur.py +52 -0
  59. stouputils/data_science/data_processing/image/noise.py +59 -0
  60. stouputils/data_science/data_processing/image/normalize.py +65 -0
  61. stouputils/data_science/data_processing/image/random_erase.py +66 -0
  62. stouputils/data_science/data_processing/image/resize.py +69 -0
  63. stouputils/data_science/data_processing/image/rotation.py +80 -0
  64. stouputils/data_science/data_processing/image/salt_pepper.py +68 -0
  65. stouputils/data_science/data_processing/image/sharpening.py +55 -0
  66. stouputils/data_science/data_processing/image/shearing.py +64 -0
  67. stouputils/data_science/data_processing/image/threshold.py +64 -0
  68. stouputils/data_science/data_processing/image/translation.py +71 -0
  69. stouputils/data_science/data_processing/image/zoom.py +83 -0
  70. stouputils/data_science/data_processing/image_augmentation.py +118 -0
  71. stouputils/data_science/data_processing/image_preprocess.py +183 -0
  72. stouputils/data_science/data_processing/prosthesis_detection.py +359 -0
  73. stouputils/data_science/data_processing/technique.py +481 -0
  74. stouputils/data_science/dataset/__init__.py +45 -0
  75. stouputils/data_science/dataset/dataset.py +292 -0
  76. stouputils/data_science/dataset/dataset_loader.py +135 -0
  77. stouputils/data_science/dataset/grouping_strategy.py +296 -0
  78. stouputils/data_science/dataset/image_loader.py +100 -0
  79. stouputils/data_science/dataset/xy_tuple.py +696 -0
  80. stouputils/data_science/metric_dictionnary.py +106 -0
  81. stouputils/data_science/metric_utils.py +847 -0
  82. stouputils/data_science/mlflow_utils.py +206 -0
  83. stouputils/data_science/models/abstract_model.py +149 -0
  84. stouputils/data_science/models/all.py +85 -0
  85. stouputils/data_science/models/base_keras.py +765 -0
  86. stouputils/data_science/models/keras/all.py +38 -0
  87. stouputils/data_science/models/keras/convnext.py +62 -0
  88. stouputils/data_science/models/keras/densenet.py +50 -0
  89. stouputils/data_science/models/keras/efficientnet.py +60 -0
  90. stouputils/data_science/models/keras/mobilenet.py +56 -0
  91. stouputils/data_science/models/keras/resnet.py +52 -0
  92. stouputils/data_science/models/keras/squeezenet.py +233 -0
  93. stouputils/data_science/models/keras/vgg.py +42 -0
  94. stouputils/data_science/models/keras/xception.py +38 -0
  95. stouputils/data_science/models/keras_utils/callbacks/__init__.py +20 -0
  96. stouputils/data_science/models/keras_utils/callbacks/colored_progress_bar.py +219 -0
  97. stouputils/data_science/models/keras_utils/callbacks/learning_rate_finder.py +148 -0
  98. stouputils/data_science/models/keras_utils/callbacks/model_checkpoint_v2.py +31 -0
  99. stouputils/data_science/models/keras_utils/callbacks/progressive_unfreezing.py +249 -0
  100. stouputils/data_science/models/keras_utils/callbacks/warmup_scheduler.py +66 -0
  101. stouputils/data_science/models/keras_utils/losses/__init__.py +12 -0
  102. stouputils/data_science/models/keras_utils/losses/next_generation_loss.py +56 -0
  103. stouputils/data_science/models/keras_utils/visualizations.py +416 -0
  104. stouputils/data_science/models/model_interface.py +939 -0
  105. stouputils/data_science/models/sandbox.py +116 -0
  106. stouputils/data_science/range_tuple.py +234 -0
  107. stouputils/data_science/scripts/augment_dataset.py +77 -0
  108. stouputils/data_science/scripts/exhaustive_process.py +133 -0
  109. stouputils/data_science/scripts/preprocess_dataset.py +70 -0
  110. stouputils/data_science/scripts/routine.py +168 -0
  111. stouputils/data_science/utils.py +285 -0
  112. stouputils/decorators.py +595 -0
  113. stouputils/decorators.pyi +242 -0
  114. stouputils/image.py +441 -0
  115. stouputils/image.pyi +172 -0
  116. stouputils/installer/__init__.py +18 -0
  117. stouputils/installer/__init__.pyi +5 -0
  118. stouputils/installer/common.py +67 -0
  119. stouputils/installer/common.pyi +39 -0
  120. stouputils/installer/downloader.py +101 -0
  121. stouputils/installer/downloader.pyi +24 -0
  122. stouputils/installer/linux.py +144 -0
  123. stouputils/installer/linux.pyi +39 -0
  124. stouputils/installer/main.py +223 -0
  125. stouputils/installer/main.pyi +57 -0
  126. stouputils/installer/windows.py +136 -0
  127. stouputils/installer/windows.pyi +31 -0
  128. stouputils/io.py +486 -0
  129. stouputils/io.pyi +213 -0
  130. stouputils/parallel.py +453 -0
  131. stouputils/parallel.pyi +211 -0
  132. stouputils/print.py +527 -0
  133. stouputils/print.pyi +146 -0
  134. stouputils/py.typed +1 -0
  135. stouputils-1.12.1.dist-info/METADATA +179 -0
  136. stouputils-1.12.1.dist-info/RECORD +138 -0
  137. stouputils-1.12.1.dist-info/WHEEL +4 -0
  138. stouputils-1.12.1.dist-info/entry_points.txt +3 -0
stouputils/parallel.py ADDED
@@ -0,0 +1,453 @@
1
+ """
2
+ This module provides utility functions for parallel processing, such as:
3
+
4
+ - multiprocessing(): Execute a function in parallel using multiprocessing
5
+ - multithreading(): Execute a function in parallel using multithreading
6
+ - run_in_subprocess(): Execute a function in a subprocess with args and kwargs
7
+
8
+ I highly encourage you to read the function docstrings to understand when to use each method.
9
+
10
+ .. image:: https://raw.githubusercontent.com/Stoupy51/stouputils/refs/heads/main/assets/parallel_module.gif
11
+ :alt: stouputils parallel examples
12
+ """
13
+
14
+ # Imports
15
+ import os
16
+ import time
17
+ from collections.abc import Callable, Iterable
18
+ from typing import Any, TypeVar, cast
19
+
20
+ from .ctx import SetMPStartMethod
21
+ from .print import BAR_FORMAT, MAGENTA
22
+
23
+
24
+ # Small test functions for doctests
25
+ def doctest_square(x: int) -> int:
26
+ return x * x
27
+ def doctest_slow(x: int) -> int:
28
+ time.sleep(0.1)
29
+ return x
30
+
31
+ # Constants
32
+ CPU_COUNT: int = cast(int, os.cpu_count())
33
+ T = TypeVar("T")
34
+ R = TypeVar("R")
35
+
36
+ # Functions
37
+ def multiprocessing[T, R](
38
+ func: Callable[..., R] | list[Callable[..., R]],
39
+ args: Iterable[T],
40
+ use_starmap: bool = False,
41
+ chunksize: int = 1,
42
+ desc: str = "",
43
+ max_workers: int | float = CPU_COUNT,
44
+ delay_first_calls: float = 0,
45
+ color: str = MAGENTA,
46
+ bar_format: str = BAR_FORMAT,
47
+ ascii: bool = False,
48
+ ) -> list[R]:
49
+ r""" Method to execute a function in parallel using multiprocessing
50
+
51
+ - For CPU-bound operations where the GIL (Global Interpreter Lock) is a bottleneck.
52
+ - When the task can be divided into smaller, independent sub-tasks that can be executed concurrently.
53
+ - For computationally intensive tasks like scientific simulations, data analysis, or machine learning workloads.
54
+
55
+ Args:
56
+ func (Callable | list[Callable]): Function to execute, or list of functions (one per argument)
57
+ args (Iterable): Iterable of arguments to pass to the function(s)
58
+ use_starmap (bool): Whether to use starmap or not (Defaults to False):
59
+ True means the function will be called like func(\*args[i]) instead of func(args[i])
60
+ chunksize (int): Number of arguments to process at a time
61
+ (Defaults to 1 for proper progress bar display)
62
+ desc (str): Description displayed in the progress bar
63
+ (if not provided no progress bar will be displayed)
64
+ max_workers (int | float): Number of workers to use (Defaults to CPU_COUNT), -1 means CPU_COUNT.
65
+ If float between 0 and 1, it's treated as a percentage of CPU_COUNT.
66
+ If negative float between -1 and 0, it's treated as a percentage of len(args).
67
+ delay_first_calls (float): Apply i*delay_first_calls seconds delay to the first "max_workers" calls.
68
+ For instance, the first process will be delayed by 0 seconds, the second by 1 second, etc.
69
+ (Defaults to 0): This can be useful to avoid functions being called in the same second.
70
+ color (str): Color of the progress bar (Defaults to MAGENTA)
71
+ bar_format (str): Format of the progress bar (Defaults to BAR_FORMAT)
72
+ ascii (bool): Whether to use ASCII or Unicode characters for the progress bar
73
+
74
+ Returns:
75
+ list[object]: Results of the function execution
76
+
77
+ Examples:
78
+ .. code-block:: python
79
+
80
+ > multiprocessing(doctest_square, args=[1, 2, 3])
81
+ [1, 4, 9]
82
+
83
+ > multiprocessing(int.__mul__, [(1,2), (3,4), (5,6)], use_starmap=True)
84
+ [2, 12, 30]
85
+
86
+ > # Using a list of functions (one per argument)
87
+ > multiprocessing([doctest_square, doctest_square, doctest_square], [1, 2, 3])
88
+ [1, 4, 9]
89
+
90
+ > # Will process in parallel with progress bar
91
+ > multiprocessing(doctest_slow, range(10), desc="Processing")
92
+ [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
93
+
94
+ > # Will process in parallel with progress bar and delay the first threads
95
+ > multiprocessing(
96
+ . doctest_slow,
97
+ . range(10),
98
+ . desc="Processing with delay",
99
+ . max_workers=2,
100
+ . delay_first_calls=0.6
101
+ . )
102
+ [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
103
+ """
104
+ # Imports
105
+ import multiprocessing as mp
106
+ from multiprocessing import Pool
107
+
108
+ from tqdm.auto import tqdm
109
+ from tqdm.contrib.concurrent import process_map # pyright: ignore[reportUnknownVariableType]
110
+
111
+ # Handle parameters
112
+ args = list(args) # Ensure we have a list (not other iterable)
113
+ if max_workers == -1:
114
+ max_workers = CPU_COUNT
115
+ if isinstance(max_workers, float):
116
+ if max_workers > 0:
117
+ assert max_workers <= 1, "max_workers as positive float must be between 0 and 1 (percentage of CPU_COUNT)"
118
+ max_workers = int(max_workers * CPU_COUNT)
119
+ else:
120
+ assert -1 <= max_workers < 0, "max_workers as negative float must be between -1 and 0 (percentage of len(args))"
121
+ max_workers = int(-max_workers * len(args))
122
+ verbose: bool = desc != ""
123
+ desc, func, args = _handle_parameters(func, args, use_starmap, delay_first_calls, max_workers, desc, color)
124
+ if bar_format == BAR_FORMAT:
125
+ bar_format = bar_format.replace(MAGENTA, color)
126
+
127
+ # Do multiprocessing only if there is more than 1 argument and more than 1 CPU
128
+ if max_workers > 1 and len(args) > 1:
129
+ def process() -> list[Any]:
130
+ if verbose:
131
+ return list(process_map(
132
+ func, args, max_workers=max_workers, chunksize=chunksize, desc=desc, bar_format=bar_format, ascii=ascii
133
+ )) # type: ignore
134
+ else:
135
+ with Pool(max_workers) as pool:
136
+ return list(pool.map(func, args, chunksize=chunksize)) # type: ignore
137
+ try:
138
+ return process()
139
+ except RuntimeError as e:
140
+ if "SemLock created in a fork context is being shared with a process in a spawn context" in str(e):
141
+
142
+ # Try with alternate start method
143
+ with SetMPStartMethod("spawn" if mp.get_start_method() != "spawn" else "fork"):
144
+ return process()
145
+ else: # Re-raise if it's not the SemLock error
146
+ raise
147
+
148
+ # Single process execution
149
+ else:
150
+ if verbose:
151
+ return [func(arg) for arg in tqdm(args, total=len(args), desc=desc, bar_format=bar_format, ascii=ascii)]
152
+ else:
153
+ return [func(arg) for arg in args]
154
+
155
+
156
+ def multithreading[T, R](
157
+ func: Callable[..., R] | list[Callable[..., R]],
158
+ args: Iterable[T],
159
+ use_starmap: bool = False,
160
+ desc: str = "",
161
+ max_workers: int | float = CPU_COUNT,
162
+ delay_first_calls: float = 0,
163
+ color: str = MAGENTA,
164
+ bar_format: str = BAR_FORMAT,
165
+ ascii: bool = False,
166
+ ) -> list[R]:
167
+ r""" Method to execute a function in parallel using multithreading, you should use it:
168
+
169
+ - For I/O-bound operations where the GIL is not a bottleneck, such as network requests or disk operations.
170
+ - When the task involves waiting for external resources, such as network responses or user input.
171
+ - For operations that involve a lot of waiting, such as GUI event handling or handling user input.
172
+
173
+ Args:
174
+ func (Callable | list[Callable]): Function to execute, or list of functions (one per argument)
175
+ args (Iterable): Iterable of arguments to pass to the function(s)
176
+ use_starmap (bool): Whether to use starmap or not (Defaults to False):
177
+ True means the function will be called like func(\*args[i]) instead of func(args[i])
178
+ desc (str): Description displayed in the progress bar
179
+ (if not provided no progress bar will be displayed)
180
+ max_workers (int | float): Number of workers to use (Defaults to CPU_COUNT), -1 means CPU_COUNT.
181
+ If float between 0 and 1, it's treated as a percentage of CPU_COUNT.
182
+ If negative float between -1 and 0, it's treated as a percentage of len(args).
183
+ delay_first_calls (float): Apply i*delay_first_calls seconds delay to the first "max_workers" calls.
184
+ For instance with value to 1, the first thread will be delayed by 0 seconds, the second by 1 second, etc.
185
+ (Defaults to 0): This can be useful to avoid functions being called in the same second.
186
+ color (str): Color of the progress bar (Defaults to MAGENTA)
187
+ bar_format (str): Format of the progress bar (Defaults to BAR_FORMAT)
188
+ ascii (bool): Whether to use ASCII or Unicode characters for the progress bar
189
+
190
+ Returns:
191
+ list[object]: Results of the function execution
192
+
193
+ Examples:
194
+ .. code-block:: python
195
+
196
+ > multithreading(doctest_square, args=[1, 2, 3])
197
+ [1, 4, 9]
198
+
199
+ > multithreading(int.__mul__, [(1,2), (3,4), (5,6)], use_starmap=True)
200
+ [2, 12, 30]
201
+
202
+ > # Using a list of functions (one per argument)
203
+ > multithreading([doctest_square, doctest_square, doctest_square], [1, 2, 3])
204
+ [1, 4, 9]
205
+
206
+ > # Will process in parallel with progress bar
207
+ > multithreading(doctest_slow, range(10), desc="Threading")
208
+ [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
209
+
210
+ > # Will process in parallel with progress bar and delay the first threads
211
+ > multithreading(
212
+ . doctest_slow,
213
+ . range(10),
214
+ . desc="Threading with delay",
215
+ . max_workers=2,
216
+ . delay_first_calls=0.6
217
+ . )
218
+ [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
219
+ """
220
+ # Imports
221
+ from concurrent.futures import ThreadPoolExecutor
222
+
223
+ from tqdm.auto import tqdm
224
+
225
+ # Handle parameters
226
+ args = list(args) # Ensure we have a list (not other iterable)
227
+ if max_workers == -1:
228
+ max_workers = CPU_COUNT
229
+ if isinstance(max_workers, float):
230
+ if max_workers > 0:
231
+ assert max_workers <= 1, "max_workers as positive float must be between 0 and 1 (percentage of CPU_COUNT)"
232
+ max_workers = int(max_workers * CPU_COUNT)
233
+ else:
234
+ assert -1 <= max_workers < 0, "max_workers as negative float must be between -1 and 0 (percentage of len(args))"
235
+ max_workers = int(-max_workers * len(args))
236
+ verbose: bool = desc != ""
237
+ desc, func, args = _handle_parameters(func, args, use_starmap, delay_first_calls, max_workers, desc, color)
238
+ if bar_format == BAR_FORMAT:
239
+ bar_format = bar_format.replace(MAGENTA, color)
240
+
241
+ # Do multithreading only if there is more than 1 argument and more than 1 CPU
242
+ if max_workers > 1 and len(args) > 1:
243
+ if verbose:
244
+ with ThreadPoolExecutor(max_workers) as executor:
245
+ return list(tqdm(executor.map(func, args), total=len(args), desc=desc, bar_format=bar_format, ascii=ascii))
246
+ else:
247
+ with ThreadPoolExecutor(max_workers) as executor:
248
+ return list(executor.map(func, args))
249
+
250
+ # Single process execution
251
+ else:
252
+ if verbose:
253
+ return [func(arg) for arg in tqdm(args, total=len(args), desc=desc, bar_format=bar_format, ascii=ascii)]
254
+ else:
255
+ return [func(arg) for arg in args]
256
+
257
+
258
+ def run_in_subprocess[R](
259
+ func: Callable[..., R],
260
+ *args: Any,
261
+ timeout: float | None = None,
262
+ **kwargs: Any
263
+ ) -> R:
264
+ """ Execute a function in a subprocess with positional and keyword arguments.
265
+
266
+ This is useful when you need to run a function in isolation to avoid memory leaks,
267
+ resource conflicts, or to ensure a clean execution environment. The subprocess will
268
+ be created, run the function with the provided arguments, and return the result.
269
+
270
+ Args:
271
+ func (Callable): The function to execute in a subprocess.
272
+ (SHOULD BE A TOP-LEVEL FUNCTION TO BE PICKLABLE)
273
+ *args (Any): Positional arguments to pass to the function.
274
+ timeout (float | None): Maximum time in seconds to wait for the subprocess.
275
+ If None, wait indefinitely. If the subprocess exceeds this time, it will be terminated.
276
+ **kwargs (Any): Keyword arguments to pass to the function.
277
+
278
+ Returns:
279
+ R: The return value of the function.
280
+
281
+ Raises:
282
+ RuntimeError: If the subprocess exits with a non-zero exit code or times out.
283
+ TimeoutError: If the subprocess exceeds the specified timeout.
284
+
285
+ Examples:
286
+ .. code-block:: python
287
+
288
+ > # Simple function execution
289
+ > run_in_subprocess(doctest_square, 5)
290
+ 25
291
+
292
+ > # Function with multiple arguments
293
+ > def add(a: int, b: int) -> int:
294
+ . return a + b
295
+ > run_in_subprocess(add, 10, 20)
296
+ 30
297
+
298
+ > # Function with keyword arguments
299
+ > def greet(name: str, greeting: str = "Hello") -> str:
300
+ . return f"{greeting}, {name}!"
301
+ > run_in_subprocess(greet, "World", greeting="Hi")
302
+ 'Hi, World!'
303
+
304
+ > # With timeout to prevent hanging
305
+ > run_in_subprocess(some_gpu_func, data, timeout=300.0)
306
+ """
307
+ import multiprocessing as mp
308
+ from multiprocessing import Queue
309
+
310
+ # Create a queue to get the result from the subprocess
311
+ result_queue: Queue[R | Exception] = Queue()
312
+
313
+ # Create and start the subprocess using the module-level wrapper
314
+ process: mp.Process = mp.Process(
315
+ target=_subprocess_wrapper,
316
+ args=(result_queue, func, args, kwargs)
317
+ )
318
+ process.start()
319
+
320
+ # Join with timeout to prevent indefinite hanging
321
+ process.join(timeout=timeout)
322
+
323
+ # Check if process is still alive (timed out)
324
+ if process.is_alive():
325
+ process.terminate()
326
+ time.sleep(0.5) # Give it a moment to terminate gracefully
327
+ if process.is_alive():
328
+ process.kill()
329
+ process.join()
330
+ raise TimeoutError(f"Subprocess exceeded timeout of {timeout} seconds and was terminated")
331
+
332
+ # Check exit code
333
+ if process.exitcode != 0:
334
+ # Try to get any exception from the queue (non-blocking)
335
+ error_msg = f"Subprocess failed with exit code {process.exitcode}"
336
+ try:
337
+ if not result_queue.empty():
338
+ result_or_exception = result_queue.get_nowait()
339
+ if isinstance(result_or_exception, Exception):
340
+ raise result_or_exception
341
+ except Exception:
342
+ pass
343
+ raise RuntimeError(error_msg)
344
+
345
+ # Retrieve the result
346
+ try:
347
+ result_or_exception = result_queue.get_nowait()
348
+ if isinstance(result_or_exception, Exception):
349
+ raise result_or_exception
350
+ return result_or_exception
351
+ except Exception as e:
352
+ raise RuntimeError("Subprocess did not return any result") from e
353
+
354
+
355
+ # "Private" function for subprocess wrapper (must be at module level for pickling on Windows)
356
+ def _subprocess_wrapper[R](
357
+ result_queue: Any,
358
+ func: Callable[..., R],
359
+ args: tuple[Any, ...],
360
+ kwargs: dict[str, Any]
361
+ ) -> None:
362
+ """ Wrapper function to execute the target function and store the result in the queue.
363
+
364
+ Must be at module level to be pickable on Windows (spawn context).
365
+
366
+ Args:
367
+ result_queue (multiprocessing.Queue): Queue to store the result or exception.
368
+ func (Callable): The target function to execute.
369
+ args (tuple): Positional arguments for the function.
370
+ kwargs (dict): Keyword arguments for the function.
371
+ """
372
+ try:
373
+ result: R = func(*args, **kwargs)
374
+ result_queue.put(result)
375
+ except Exception as e:
376
+ result_queue.put(e)
377
+
378
+ # "Private" function to use starmap
379
+ def _starmap[T, R](args: tuple[Callable[[T], R], list[T]]) -> R:
380
+ r""" Private function to use starmap using args[0](\*args[1])
381
+
382
+ Args:
383
+ args (tuple): Tuple containing the function and the arguments list to pass to the function
384
+ Returns:
385
+ object: Result of the function execution
386
+ """
387
+ func, arguments = args
388
+ return func(*arguments)
389
+
390
+ # "Private" function to apply delay before calling the target function
391
+ def _delayed_call[T, R](args: tuple[Callable[[T], R], float, T]) -> R:
392
+ """ Private function to apply delay before calling the target function
393
+
394
+ Args:
395
+ args (tuple): Tuple containing the function, delay in seconds, and the argument to pass to the function
396
+ Returns:
397
+ object: Result of the function execution
398
+ """
399
+ func, delay, arg = args
400
+ time.sleep(delay)
401
+ return func(arg)
402
+
403
+ # "Private" function to handle parameters for multiprocessing or multithreading functions
404
+ def _handle_parameters[T, R](
405
+ func: Callable[[T], R] | list[Callable[[T], R]],
406
+ args: list[T],
407
+ use_starmap: bool,
408
+ delay_first_calls: float,
409
+ max_workers: int,
410
+ desc: str,
411
+ color: str
412
+ ) -> tuple[str, Callable[[T], R], list[T]]:
413
+ r""" Private function to handle the parameters for multiprocessing or multithreading functions
414
+
415
+ Args:
416
+ func (Callable | list[Callable]): Function to execute, or list of functions (one per argument)
417
+ args (list): List of arguments to pass to the function(s)
418
+ use_starmap (bool): Whether to use starmap or not (Defaults to False):
419
+ True means the function will be called like func(\*args[i]) instead of func(args[i])
420
+ delay_first_calls (int): Apply i*delay_first_calls seconds delay to the first "max_workers" calls.
421
+ For instance, the first process will be delayed by 0 seconds, the second by 1 second, etc. (Defaults to 0):
422
+ This can be useful to avoid functions being called in the same second.
423
+ max_workers (int): Number of workers to use (Defaults to CPU_COUNT)
424
+ desc (str): Description of the function execution displayed in the progress bar
425
+ color (str): Color of the progress bar
426
+
427
+ Returns:
428
+ tuple[str, Callable[[T], R], list[T]]: Tuple containing the description, function, and arguments
429
+ """
430
+ desc = color + desc
431
+
432
+ # Handle list of functions: validate and convert to starmap format
433
+ if isinstance(func, list):
434
+ func = cast(list[Callable[[T], R]], func)
435
+ assert len(func) == len(args), f"Length mismatch: {len(func)} functions but {len(args)} arguments"
436
+ args = [(f, arg if use_starmap else (arg,)) for f, arg in zip(func, args, strict=False)] # type: ignore
437
+ func = _starmap # type: ignore
438
+
439
+ # If use_starmap is True, we use the _starmap function
440
+ elif use_starmap:
441
+ args = [(func, arg) for arg in args] # type: ignore
442
+ func = _starmap # type: ignore
443
+
444
+ # Prepare delayed function calls if delay_first_calls is set
445
+ if delay_first_calls > 0:
446
+ args = [
447
+ (func, i * delay_first_calls if i < max_workers else 0, arg) # type: ignore
448
+ for i, arg in enumerate(args)
449
+ ]
450
+ func = _delayed_call # type: ignore
451
+
452
+ return desc, func, args # type: ignore
453
+
@@ -0,0 +1,211 @@
1
+ from .ctx import SetMPStartMethod as SetMPStartMethod
2
+ from .print import BAR_FORMAT as BAR_FORMAT, MAGENTA as MAGENTA
3
+ from collections.abc import Callable, Iterable
4
+ from typing import Any, TypeVar
5
+
6
+ def doctest_square(x: int) -> int: ...
7
+ def doctest_slow(x: int) -> int: ...
8
+
9
+ CPU_COUNT: int
10
+ T = TypeVar('T')
11
+ R = TypeVar('R')
12
+
13
+ def multiprocessing[T, R](func: Callable[..., R] | list[Callable[..., R]], args: Iterable[T], use_starmap: bool = False, chunksize: int = 1, desc: str = '', max_workers: int | float = ..., delay_first_calls: float = 0, color: str = ..., bar_format: str = ..., ascii: bool = False) -> list[R]:
14
+ ''' Method to execute a function in parallel using multiprocessing
15
+
16
+ \t- For CPU-bound operations where the GIL (Global Interpreter Lock) is a bottleneck.
17
+ \t- When the task can be divided into smaller, independent sub-tasks that can be executed concurrently.
18
+ \t- For computationally intensive tasks like scientific simulations, data analysis, or machine learning workloads.
19
+
20
+ \tArgs:
21
+ \t\tfunc\t\t\t\t(Callable | list[Callable]):\tFunction to execute, or list of functions (one per argument)
22
+ \t\targs\t\t\t\t(Iterable):\t\t\tIterable of arguments to pass to the function(s)
23
+ \t\tuse_starmap\t\t\t(bool):\t\t\t\tWhether to use starmap or not (Defaults to False):
24
+ \t\t\tTrue means the function will be called like func(\\*args[i]) instead of func(args[i])
25
+ \t\tchunksize\t\t\t(int):\t\t\t\tNumber of arguments to process at a time
26
+ \t\t\t(Defaults to 1 for proper progress bar display)
27
+ \t\tdesc\t\t\t\t(str):\t\t\t\tDescription displayed in the progress bar
28
+ \t\t\t(if not provided no progress bar will be displayed)
29
+ \t\tmax_workers\t\t\t(int | float):\t\tNumber of workers to use (Defaults to CPU_COUNT), -1 means CPU_COUNT.
30
+ \t\t\tIf float between 0 and 1, it\'s treated as a percentage of CPU_COUNT.
31
+ \t\t\tIf negative float between -1 and 0, it\'s treated as a percentage of len(args).
32
+ \t\tdelay_first_calls\t(float):\t\t\tApply i*delay_first_calls seconds delay to the first "max_workers" calls.
33
+ \t\t\tFor instance, the first process will be delayed by 0 seconds, the second by 1 second, etc.
34
+ \t\t\t(Defaults to 0): This can be useful to avoid functions being called in the same second.
35
+ \t\tcolor\t\t\t\t(str):\t\t\t\tColor of the progress bar (Defaults to MAGENTA)
36
+ \t\tbar_format\t\t\t(str):\t\t\t\tFormat of the progress bar (Defaults to BAR_FORMAT)
37
+ \t\tascii\t\t\t\t(bool):\t\t\t\tWhether to use ASCII or Unicode characters for the progress bar
38
+
39
+ \tReturns:
40
+ \t\tlist[object]:\tResults of the function execution
41
+
42
+ \tExamples:
43
+ \t\t.. code-block:: python
44
+
45
+ \t\t\t> multiprocessing(doctest_square, args=[1, 2, 3])
46
+ \t\t\t[1, 4, 9]
47
+
48
+ \t\t\t> multiprocessing(int.__mul__, [(1,2), (3,4), (5,6)], use_starmap=True)
49
+ \t\t\t[2, 12, 30]
50
+
51
+ \t\t\t> # Using a list of functions (one per argument)
52
+ \t\t\t> multiprocessing([doctest_square, doctest_square, doctest_square], [1, 2, 3])
53
+ \t\t\t[1, 4, 9]
54
+
55
+ \t\t\t> # Will process in parallel with progress bar
56
+ \t\t\t> multiprocessing(doctest_slow, range(10), desc="Processing")
57
+ \t\t\t[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
58
+
59
+ \t\t\t> # Will process in parallel with progress bar and delay the first threads
60
+ \t\t\t> multiprocessing(
61
+ \t\t\t. doctest_slow,
62
+ \t\t\t. range(10),
63
+ \t\t\t. desc="Processing with delay",
64
+ \t\t\t. max_workers=2,
65
+ \t\t\t. delay_first_calls=0.6
66
+ \t\t\t. )
67
+ \t\t\t[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
68
+ \t'''
69
+ def multithreading[T, R](func: Callable[..., R] | list[Callable[..., R]], args: Iterable[T], use_starmap: bool = False, desc: str = '', max_workers: int | float = ..., delay_first_calls: float = 0, color: str = ..., bar_format: str = ..., ascii: bool = False) -> list[R]:
70
+ ''' Method to execute a function in parallel using multithreading, you should use it:
71
+
72
+ \t- For I/O-bound operations where the GIL is not a bottleneck, such as network requests or disk operations.
73
+ \t- When the task involves waiting for external resources, such as network responses or user input.
74
+ \t- For operations that involve a lot of waiting, such as GUI event handling or handling user input.
75
+
76
+ \tArgs:
77
+ \t\tfunc\t\t\t\t(Callable | list[Callable]):\tFunction to execute, or list of functions (one per argument)
78
+ \t\targs\t\t\t\t(Iterable):\t\t\tIterable of arguments to pass to the function(s)
79
+ \t\tuse_starmap\t\t\t(bool):\t\t\t\tWhether to use starmap or not (Defaults to False):
80
+ \t\t\tTrue means the function will be called like func(\\*args[i]) instead of func(args[i])
81
+ \t\tdesc\t\t\t\t(str):\t\t\t\tDescription displayed in the progress bar
82
+ \t\t\t(if not provided no progress bar will be displayed)
83
+ \t\tmax_workers\t\t\t(int | float):\t\tNumber of workers to use (Defaults to CPU_COUNT), -1 means CPU_COUNT.
84
+ \t\t\tIf float between 0 and 1, it\'s treated as a percentage of CPU_COUNT.
85
+ \t\t\tIf negative float between -1 and 0, it\'s treated as a percentage of len(args).
86
+ \t\tdelay_first_calls\t(float):\t\t\tApply i*delay_first_calls seconds delay to the first "max_workers" calls.
87
+ \t\t\tFor instance with value to 1, the first thread will be delayed by 0 seconds, the second by 1 second, etc.
88
+ \t\t\t(Defaults to 0): This can be useful to avoid functions being called in the same second.
89
+ \t\tcolor\t\t\t\t(str):\t\t\t\tColor of the progress bar (Defaults to MAGENTA)
90
+ \t\tbar_format\t\t\t(str):\t\t\t\tFormat of the progress bar (Defaults to BAR_FORMAT)
91
+ \t\tascii\t\t\t\t(bool):\t\t\t\tWhether to use ASCII or Unicode characters for the progress bar
92
+
93
+ \tReturns:
94
+ \t\tlist[object]:\tResults of the function execution
95
+
96
+ \tExamples:
97
+ \t\t.. code-block:: python
98
+
99
+ \t\t\t> multithreading(doctest_square, args=[1, 2, 3])
100
+ \t\t\t[1, 4, 9]
101
+
102
+ \t\t\t> multithreading(int.__mul__, [(1,2), (3,4), (5,6)], use_starmap=True)
103
+ \t\t\t[2, 12, 30]
104
+
105
+ \t\t\t> # Using a list of functions (one per argument)
106
+ \t\t\t> multithreading([doctest_square, doctest_square, doctest_square], [1, 2, 3])
107
+ \t\t\t[1, 4, 9]
108
+
109
+ \t\t\t> # Will process in parallel with progress bar
110
+ \t\t\t> multithreading(doctest_slow, range(10), desc="Threading")
111
+ \t\t\t[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
112
+
113
+ \t\t\t> # Will process in parallel with progress bar and delay the first threads
114
+ \t\t\t> multithreading(
115
+ \t\t\t. doctest_slow,
116
+ \t\t\t. range(10),
117
+ \t\t\t. desc="Threading with delay",
118
+ \t\t\t. max_workers=2,
119
+ \t\t\t. delay_first_calls=0.6
120
+ \t\t\t. )
121
+ \t\t\t[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
122
+ \t'''
123
+ def run_in_subprocess[R](func: Callable[..., R], *args: Any, timeout: float | None = None, **kwargs: Any) -> R:
124
+ ''' Execute a function in a subprocess with positional and keyword arguments.
125
+
126
+ \tThis is useful when you need to run a function in isolation to avoid memory leaks,
127
+ \tresource conflicts, or to ensure a clean execution environment. The subprocess will
128
+ \tbe created, run the function with the provided arguments, and return the result.
129
+
130
+ \tArgs:
131
+ \t\tfunc (Callable): The function to execute in a subprocess.
132
+ \t\t\t(SHOULD BE A TOP-LEVEL FUNCTION TO BE PICKLABLE)
133
+ \t\t*args (Any): Positional arguments to pass to the function.
134
+ \t\ttimeout (float | None): Maximum time in seconds to wait for the subprocess.
135
+ \t\t\tIf None, wait indefinitely. If the subprocess exceeds this time, it will be terminated.
136
+ \t\t**kwargs (Any): Keyword arguments to pass to the function.
137
+
138
+ \tReturns:
139
+ \t\tR: The return value of the function.
140
+
141
+ \tRaises:
142
+ \t\tRuntimeError: If the subprocess exits with a non-zero exit code or times out.
143
+ \t\tTimeoutError: If the subprocess exceeds the specified timeout.
144
+
145
+ \tExamples:
146
+ \t\t.. code-block:: python
147
+
148
+ \t\t\t> # Simple function execution
149
+ \t\t\t> run_in_subprocess(doctest_square, 5)
150
+ \t\t\t25
151
+
152
+ \t\t\t> # Function with multiple arguments
153
+ \t\t\t> def add(a: int, b: int) -> int:
154
+ \t\t\t. return a + b
155
+ \t\t\t> run_in_subprocess(add, 10, 20)
156
+ \t\t\t30
157
+
158
+ \t\t\t> # Function with keyword arguments
159
+ \t\t\t> def greet(name: str, greeting: str = "Hello") -> str:
160
+ \t\t\t. return f"{greeting}, {name}!"
161
+ \t\t\t> run_in_subprocess(greet, "World", greeting="Hi")
162
+ \t\t\t\'Hi, World!\'
163
+
164
+ \t\t\t> # With timeout to prevent hanging
165
+ \t\t\t> run_in_subprocess(some_gpu_func, data, timeout=300.0)
166
+ \t'''
167
+ def _subprocess_wrapper[R](result_queue: Any, func: Callable[..., R], args: tuple[Any, ...], kwargs: dict[str, Any]) -> None:
168
+ """ Wrapper function to execute the target function and store the result in the queue.
169
+
170
+ \tMust be at module level to be pickable on Windows (spawn context).
171
+
172
+ \tArgs:
173
+ \t\tresult_queue (multiprocessing.Queue): Queue to store the result or exception.
174
+ \t\tfunc (Callable): The target function to execute.
175
+ \t\targs (tuple): Positional arguments for the function.
176
+ \t\tkwargs (dict): Keyword arguments for the function.
177
+ \t"""
178
+ def _starmap[T, R](args: tuple[Callable[[T], R], list[T]]) -> R:
179
+ """ Private function to use starmap using args[0](\\*args[1])
180
+
181
+ \tArgs:
182
+ \t\targs (tuple): Tuple containing the function and the arguments list to pass to the function
183
+ \tReturns:
184
+ \t\tobject: Result of the function execution
185
+ \t"""
186
+ def _delayed_call[T, R](args: tuple[Callable[[T], R], float, T]) -> R:
187
+ """ Private function to apply delay before calling the target function
188
+
189
+ \tArgs:
190
+ \t\targs (tuple): Tuple containing the function, delay in seconds, and the argument to pass to the function
191
+ \tReturns:
192
+ \t\tobject: Result of the function execution
193
+ \t"""
194
+ def _handle_parameters[T, R](func: Callable[[T], R] | list[Callable[[T], R]], args: list[T], use_starmap: bool, delay_first_calls: float, max_workers: int, desc: str, color: str) -> tuple[str, Callable[[T], R], list[T]]:
195
+ ''' Private function to handle the parameters for multiprocessing or multithreading functions
196
+
197
+ \tArgs:
198
+ \t\tfunc\t\t\t\t(Callable | list[Callable]):\tFunction to execute, or list of functions (one per argument)
199
+ \t\targs\t\t\t\t(list):\t\t\t\tList of arguments to pass to the function(s)
200
+ \t\tuse_starmap\t\t\t(bool):\t\t\t\tWhether to use starmap or not (Defaults to False):
201
+ \t\t\tTrue means the function will be called like func(\\*args[i]) instead of func(args[i])
202
+ \t\tdelay_first_calls\t(int):\t\t\t\tApply i*delay_first_calls seconds delay to the first "max_workers" calls.
203
+ \t\t\tFor instance, the first process will be delayed by 0 seconds, the second by 1 second, etc. (Defaults to 0):
204
+ \t\t\tThis can be useful to avoid functions being called in the same second.
205
+ \t\tmax_workers\t\t\t(int):\t\t\t\tNumber of workers to use (Defaults to CPU_COUNT)
206
+ \t\tdesc\t\t\t\t(str):\t\t\t\tDescription of the function execution displayed in the progress bar
207
+ \t\tcolor\t\t\t\t(str):\t\t\t\tColor of the progress bar
208
+
209
+ \tReturns:
210
+ \t\ttuple[str, Callable[[T], R], list[T]]:\tTuple containing the description, function, and arguments
211
+ \t'''