forked from SQLMesh/sqlmesh
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathprocess.py
More file actions
73 lines (57 loc) · 2.37 KB
/
Copy pathprocess.py
File metadata and controls
73 lines (57 loc) · 2.37 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
# mypy: disable-error-code=no-untyped-def
from concurrent.futures import Future, ProcessPoolExecutor
import typing as t
import multiprocessing as mp
from sqlmesh.utils.windows import IS_WINDOWS
class SynchronousPoolExecutor:
"""A mock implementation of the ProcessPoolExecutor for synchronous use.
This executor runs functions synchronously in the same process, avoiding the issues
with forking in test environments or when forking isn't possible (non-posix).
"""
def __init__(self, max_workers=None, mp_context=None, initializer=None, initargs=()):
if initializer is not None:
try:
initializer(*initargs)
except BaseException as ex:
raise RuntimeError(f"Exception in initializer: {ex}")
def __enter__(self):
return self
def __exit__(self, *args):
self.shutdown(wait=True)
return False
def shutdown(self, wait=True, cancel_futures=False):
"""No-op method to match ProcessPoolExecutor API.
Since this executor runs synchronously, there are no background processes
or resources to shut down and all futures will have completed already.
"""
pass
def submit(self, fn, *args, **kwargs):
"""Execute the function synchronously and return a Future with the result."""
future = Future()
try:
result = fn(*args, **kwargs)
future.set_result(result)
except Exception as e:
future.set_exception(e)
return future
def map(self, fn, *iterables, timeout=None, chunksize=1):
"""Synchronous implementation of ProcessPoolExecutor.map.
This executes the function for each set of inputs from the iterables in the
current process using Python's built-in map, rather than distributing work.
"""
return map(fn, *iterables)
PoolExecutor = t.Union[SynchronousPoolExecutor, ProcessPoolExecutor]
def create_process_pool_executor(
initializer: t.Callable, initargs: t.Tuple, max_workers: t.Optional[int]
) -> PoolExecutor:
if max_workers == 1 or IS_WINDOWS:
return SynchronousPoolExecutor(
initializer=initializer,
initargs=initargs,
)
return ProcessPoolExecutor(
mp_context=mp.get_context("fork"),
initializer=initializer,
initargs=initargs,
max_workers=max_workers,
)