1
1
import subprocess
2
2
import os
3
+ import queue
4
+ import platform
3
5
from concurrent .futures import ThreadPoolExecutor
4
6
from multiprocessing import Event
5
7
from tqdm import tqdm
6
8
7
- import platform
9
+ from Interlace . lib . core . output import OutputHelper , Level
8
10
9
11
if platform .system ().lower () == 'linux' :
10
12
shell = os .getenv ("SHELL" ) if os .getenv ("SHELL" ) else "/bin/sh"
@@ -60,8 +62,7 @@ def _run_task(self, t=False):
60
62
stdout = subprocess .DEVNULL ,
61
63
encoding = "utf-8" ,
62
64
executable = shell )
63
- out , _ = s .communicate ()
64
-
65
+ s .communicate ()
65
66
return
66
67
else :
67
68
s = subprocess .Popen (self .task , shell = True ,
@@ -78,57 +79,54 @@ def _run_task(self, t=False):
78
79
79
80
80
81
class Worker (object ):
81
- def __init__ (self , task_queue , timeout , output , tq ):
82
+ def __init__ (self , task_queue , timeout , output , tq , output_helper ):
82
83
self .queue = task_queue
83
84
self .timeout = timeout
84
85
self .output = output
85
86
self .tqdm = tq
87
+ self .output_helper = output_helper
86
88
87
89
def __call__ (self ):
88
- queue = self .queue
89
90
while True :
90
91
try :
91
- task = next (queue )
92
+ task = self .queue .get (timeout = 1 )
93
+ except queue .Empty :
94
+ return
95
+
96
+ self .output_helper .terminal (Level .THREAD , task .name (), "Added to Queue" )
97
+
98
+ try :
92
99
if isinstance (self .tqdm , tqdm ):
93
100
self .tqdm .update (1 )
94
- # run task
95
101
task .run (self .tqdm )
96
102
else :
97
103
task .run ()
98
- except StopIteration :
99
- break
104
+ except Exception as e :
105
+ self . output_helper . terminal ( Level . ERROR , task . name (), f"Task failed: { e } " )
100
106
101
107
102
108
class Pool (object ):
103
- def __init__ (self , max_workers , task_queue , timeout , output , progress_bar , silent = False ):
104
-
105
- # convert stdin input to integer
109
+ def __init__ (self , max_workers , task_queue , timeout , output , progress_bar , silent = False , output_helper = None ):
106
110
max_workers = int (max_workers )
107
-
108
- # check if there are enough workers
109
- if max_workers <= 0 :
110
- raise ValueError ("Workers must be >= 1" )
111
-
112
111
tasks_count = next (task_queue )
113
-
114
- # check if the queue is empty
115
112
if not tasks_count :
116
113
raise ValueError ("The queue is empty" )
117
114
118
- self .queue = task_queue
115
+ self .queue = queue .Queue ()
116
+ for task in task_queue :
117
+ self .queue .put (task )
118
+
119
119
self .timeout = timeout
120
120
self .output = output
121
121
self .max_workers = min (tasks_count , max_workers )
122
122
123
- if not progress_bar and not silent :
124
- self .tqdm = tqdm (total = tasks_count )
125
- else :
126
- self .tqdm = True
123
+ self .output_helper = output_helper or OutputHelper ()
124
+ self .tqdm = tqdm (total = tasks_count ) if not progress_bar and not silent else True
127
125
128
126
def run (self ):
129
- workers = [Worker (self .queue , self .timeout , self .output , self .tqdm ) for w in range (self .max_workers )]
127
+ workers = [Worker (self .queue , self .timeout , self .output , self .tqdm , self .output_helper )
128
+ for _ in range (self .max_workers )]
130
129
131
- # run
132
130
with ThreadPoolExecutor (self .max_workers ) as executors :
133
131
for worker in workers :
134
132
executors .submit (worker )
@@ -147,5 +145,5 @@ def run(self):
147
145
"sleep 9" ,
148
146
"sleep 1" ,
149
147
"echo 'Char!'" ]
150
- p = Pool (4 , tasks , 0 , 0 , True )
148
+ p = Pool (4 , iter ([ len ( tasks )] + [ Task ( t ) for t in tasks ]) , 0 , 0 , True , output_helper = OutputHelper () )
151
149
p .run ()
0 commit comments