<commit_before>import threading
import collections
import contextlib


NO_RESULT = object()


class Task:

    def __init__(self):
        self._done = False
        self._result = NO_RESULT

    def done(self, result=NO_RESULT):
        self._result = result
        self._done = True

    def is_done(self):
        return self._done

    def run(self, connection):
        raise NotImplementedError

    def get_result(self):
        return self._result


class ProcessActive(Exception):
    pass


class TaskManager:

    def __init__(self):
        self.process_active = False
        self.task_iterator = None
        self.next_task = None
        self.return_queue = None
        self.result_queue = None
        self.unfinished_tasks = 0
        self.cancelled = False

        self.mutex = threading.Lock()
        self.tasks_available = threading.Condition(self.mutex)
        self.results_available = threading.Condition(self.mutex)

    def process(self, task_iterable):
        with self.mutex:
            if self.process_active:
                raise ProcessActive
            self.process_active = True

        self.task_iterator = iter(task_iterable)
        try:
            self.next_task = next(task_iterator)
        except StopIteration:
            self.task_iterator = None
            return
        self.return_queue = collections.deque()
        self.result_queue = collections.deque()
        self.unfinished_tasks = 1
        self.cancelled = False

        with self.mutex:
            self.tasks_available.notify()

        while True:
            with self.mutex:
                if len(self.result_queue) == 0 and self.unfinished_tasks == 0:
                    break
                while len(self.result_queue) == 0:
                    self.results_available.wait()
                result = self.result_queue.pop()
            if result is not NO_RESULT:
                yield result

        self.task_iterator = None
        self.return_queue = None
        self.result_queue = None
        self.cancelled = False

        with self.mutex:
            self.process_active = False

    def cancel_process(self):
        with self.mutex:
            if self.process_active:
                self.unfinished_tasks -= len(self.return_queue)
                self.return_queue.clear()
                if self.next_task is not None:
                    self.next_task = None
                    self.unfinished_tasks -= 1
                self.cancelled = True

    def _get_next_task(self):
        with self.mutex:
            while self.next_task is None and len(self.return_queue) == 0:
                self.tasks_available.wait()

            if len(self.return_queue) > 0:
                task = self.return_queue.pop()
            else:
                task = self.next_task
                if self.cancelled:
                    self.next_task = None
                else:
                    try:
                        self.next_task = next(self.task_iterator)
                        self.unfinished_tasks += 1
                        self.tasks_available.notify()
                    except StopIteration:
                        self.next_task = None

            return task

    @contextlib.contextmanager
    def handle_task(self):
        task = self._get_next_task()
        try:
            yield task
        finally:
            with self.mutex:
                self.unfinished_tasks -= 1
                if task.is_done():
                    self.result_queue.append(task.get_result())
                    self.results_available.notify()
                elif not self.cancelled:
                    self.return_queue.append(task)
                    self.unfinished_tasks += 1
                    self.tasks_available.notify()
<commit_msg>separated task queue from main task manager code<commit_after>import threading
import collections
import contextlib


NO_RESULT = object()


class Task:

    def __init__(self):
        self._done = False
        self._result = NO_RESULT

    def done(self, result=NO_RESULT):
        self._result = result
        self._done = True

    def is_done(self):
        return self._done

    def run(self, connection):
        raise NotImplementedError

    def get_result(self):
        return self._result


class ProcessActive(Exception):
    pass


class Empty(Exception):
    pass


class TaskQueue:

    def __init__(self, task_iterable):
        self.task_iterator = iter(task_iterable)
        try:
            self.next_task = next(self.task_iterator)
            self.unfinished_tasks = 1
        except StopIteration:
            self.next_task = None
            self.unfinished_tasks = 0
        self.return_queue = collections.deque()

    def next(self):
        if len(self.return_queue) > 0:
            return self.return_queue.pop()
        elif self.next_task is not None:
            task = self.next_task
            try:
                self.next_task = next(self.task_iterator)
                self.unfinished_tasks += 1
            except StopIteration:
                self.next_task = None
            return task
        else:
            raise Empty

    def has_next(self):
        return len(self.return_queue) > 0 or self.next_task is not None

    def clear(self):
        if self.next_task is not None:
            self.next_task = None
            self.unfinished_tasks -= 1
        self.unfinished_tasks -= len(self.return_queue)
        self.return_queue.clear()

    def task_done(self):
        self.unfinished_tasks -= 1

    def has_unfinished_tasks(self):
        return self.unfinished_tasks > 0

    def return_task(self, task):
        self.return_queue.append(task)


class TaskManager:

    def __init__(self):
        self.process_active = False
        self.task_queue = None
        self.result_queue = None
        self.cancelled = False

        self.mutex = threading.Lock()
        self.tasks_available = threading.Condition(self.mutex)
        self.results_available = threading.Condition(self.mutex)

    def process(self, task_iterable):
        try:

            with self.mutex:
                if self.process_active:
                    raise ProcessActive
                self.process_active = True

            self.task_queue = TaskQueue(task_iterable)
            if not self.task_queue.has_next():
                self.task_queue = None
                return
            self.result_queue = collections.deque()
            self.cancelled = False

            with self.mutex:
                self.tasks_available.notify()

            while True:
                with self.mutex:
                    if len(self.result_queue) == 0 and not self.task_queue.has_unfinished_tasks():
                        break
                    while len(self.result_queue) == 0:
                        self.results_available.wait()
                    result = self.result_queue.pop()
                if result is not NO_RESULT:
                    yield result

        finally:

            self.task_queue = None
            self.result_queue = None
            self.cancelled = False

            with self.mutex:
                self.process_active = False

    def cancel_process(self):
        with self.mutex:
            if self.process_active:
                self.task_queue.clear()
                self.cancelled = True

    def _get_next_task(self):
        with self.mutex:
            while self.task_queue is None or not self.task_queue.has_next():
                self.tasks_available.wait()

            task = self.task_queue.next()
            if self.task_queue.has_next():
                self.tasks_available.notify()
            return task

    @contextlib.contextmanager
    def handle_task(self):
        task = self._get_next_task()
        try:
            yield task
        finally:
            with self.mutex:
                if task.is_done():
                    self.task_queue.task_done()
                    self.result_queue.append(task.get_result())
                    self.results_available.notify()
                elif not self.cancelled:
                    self.task_queue.return_task(task)
                    self.tasks_available.notify()
                else:
                    self.task_queue.task_done()

{'fraction_non_alphanumeric': 0.07374382455717557, 'fraction_numerical': 0.002530425352452103, 'mean_word_length': 1.5444512568976088, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4284952', 'n_tokens_mistral': 2257, 'n_tokens_neox': 2134, 'n_words': 477}