[ext] transmit.py file transmitter 修改完成

This commit is contained in:
li.zhenye 2022-08-11 18:21:43 +08:00
parent c076e04f57
commit 631b46d99a

View File

@ -74,6 +74,11 @@ class Transmitter(object):
self._thread_stop.set() self._thread_stop.set()
self._running_handler = None self._running_handler = None
def __del__(self):
self.stop()
if self._running_handler is not None:
self._running_handler.join()
@staticmethod @staticmethod
def job_decorator(func): def job_decorator(func):
functools.wraps(func) functools.wraps(func)
@ -128,10 +133,11 @@ class FileReceiver(Transmitter):
self.name_pattern = name_pattern self.name_pattern = name_pattern
self.file_idx = 0 self.file_idx = 0
self.output_queue = None self.output_queue = None
self.preprocess_method = None
self.set_source(input_dir, name_pattern) self.set_source(input_dir, name_pattern)
self.set_output(output_queue) self.set_output(output_queue)
def set_source(self, input_dir, name_pattern=None): def set_source(self, input_dir:str, name_pattern:str=None, preprocess_method:callable=None):
self.name_pattern = name_pattern if name_pattern is not None else self.name_pattern self.name_pattern = name_pattern if name_pattern is not None else self.name_pattern
file_names = os.listdir(input_dir) file_names = os.listdir(input_dir)
if len(file_names) == 0: if len(file_names) == 0:
@ -153,8 +159,16 @@ class FileReceiver(Transmitter):
def job_func(self, *args, **kwargs): def job_func(self, *args, **kwargs):
with self._io_lock: with self._io_lock:
self.file_idx += 1 self.file_idx += 1
if self.file_idx == len() if self.file_idx >= len(self.file_names):
self.file_idx = 0
file_name = self.file_names[self.file_idx] file_name = self.file_names[self.file_idx]
file_name = os.path.join(self.input_dir, file_name)
with open(file_name, 'rb') as f:
data = f.read()
if self.preprocess_method is not None:
data = self.preprocess_method(data)
self.output_queue.put(data)
class FifoReceiver(Transmitter): class FifoReceiver(Transmitter):