|
from time import perf_counter |
|
import logging |
|
from rich.console import Console |
|
|
|
logger = logging.getLogger(__name__) |
|
console = Console() |
|
|
|
|
|
class BaseHandler: |
|
""" |
|
Base class for pipeline parts. Each part of the pipeline has an input and an output queue. |
|
The `setup` method along with `setup_args` and `setup_kwargs` can be used to address the specific requirements of the implemented pipeline part. |
|
To stop a handler properly, set the stop_event and, to avoid queue deadlocks, place b"END" in the input queue. |
|
Objects placed in the input queue will be processed by the `process` method, and the yielded results will be placed in the output queue. |
|
The cleanup method handles stopping the handler, and b"END" is placed in the output queue. |
|
""" |
|
|
|
def __init__(self, stop_event, queue_in, queue_out, setup_args=(), setup_kwargs={}): |
|
self.stop_event = stop_event |
|
self.queue_in = queue_in |
|
self.queue_out = queue_out |
|
self.setup(*setup_args, **setup_kwargs) |
|
self._times = [] |
|
|
|
def setup(self): |
|
pass |
|
|
|
def process(self): |
|
raise NotImplementedError |
|
|
|
def run(self): |
|
while not self.stop_event.is_set(): |
|
input = self.queue_in.get() |
|
if isinstance(input, bytes) and input == b"END": |
|
|
|
logger.debug("Stopping thread") |
|
break |
|
start_time = perf_counter() |
|
for output in self.process(input): |
|
self._times.append(perf_counter() - start_time) |
|
if self.last_time > self.min_time_to_debug: |
|
console.print(f"{self.__class__.__name__}: {self.last_time: .3f} s") |
|
logger.debug(f"{self.__class__.__name__}: {self.last_time: .3f} s") |
|
self.queue_out.put(output) |
|
start_time = perf_counter() |
|
|
|
self.cleanup() |
|
self.queue_out.put(b"END") |
|
|
|
@property |
|
def last_time(self): |
|
return self._times[-1] |
|
|
|
@property |
|
def min_time_to_debug(self): |
|
return 0.001 |
|
|
|
def cleanup(self): |
|
pass |
|
|
