<commit_before># Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

import threading
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Callable, Iterable, Iterator, Optional, Sequence, Tuple

from pants.engine.fs import Digest
from pants.engine.internals.scheduler import SchedulerSession
from pants.util.logging import LogLevel


@dataclass(frozen=True)
class StreamingWorkunitContext:
    _scheduler: SchedulerSession

    def digests_to_bytes(self, digests: Sequence[Digest]) -> Tuple[bytes]:
        """Given a list of Digest objects, return a list of the bytes corresponding to each of those
        Digests in sequence."""
        return self._scheduler.digests_to_bytes(digests)

    def ensure_remote_has_recursive(self, digests: Sequence[Digest]) -> None:
        """Invoke the internal ensure_remote_has_recursive function, which ensures that a remote
        ByteStore, if it exists, has a copy of the files fingerprinted by each Digest."""
        return self._scheduler.ensure_remote_has_recursive(digests)


class StreamingWorkunitHandler:
    """StreamingWorkunitHandler's job is to periodically call each registered callback function with
    the following kwargs:

    workunits: Tuple[Dict[str, str], ...] - the workunit data itself
    finished: bool - this will be set to True when the last chunk of workunit data is reported to the callback
    """

    def __init__(
        self,
        scheduler: Any,
        callbacks: Iterable[Callable],
        report_interval_seconds: float,
        max_workunit_verbosity: LogLevel = LogLevel.DEBUG,
    ):
        self.scheduler = scheduler
        self.report_interval = report_interval_seconds
        self.callbacks = callbacks
        self._thread_runner: Optional[_InnerHandler] = None
        # TODO(10092) The max verbosity should be a per-client setting, rather than a global setting.
        self.max_workunit_verbosity = max_workunit_verbosity

    def start(self) -> None:
        if self.callbacks:
            self._thread_runner = _InnerHandler(
                self.scheduler, self.callbacks, self.report_interval, self.max_workunit_verbosity
            )
            self._thread_runner.start()

    def end(self) -> None:
        if self._thread_runner:
            self._thread_runner.join()

        # After stopping the thread, poll workunits one last time to make sure
        # we report any workunits that were added after the last time the thread polled.
        workunits = self.scheduler.poll_workunits(self.max_workunit_verbosity)
        for callback in self.callbacks:
            callback(
                workunits=workunits["completed"],
                started_workunits=workunits["started"],
                finished=True,
                context=StreamingWorkunitContext(_scheduler=self.scheduler),
            )

    @contextmanager
    def session(self) -> Iterator[None]:
        try:
            self.start()
            yield
            self.end()
        except Exception as e:
            if self._thread_runner:
                self._thread_runner.join()
            raise e


class _InnerHandler(threading.Thread):
    def __init__(
        self,
        scheduler: Any,
        callbacks: Iterable[Callable],
        report_interval: float,
        max_workunit_verbosity: LogLevel,
    ):
        super().__init__(daemon=True)
        self.scheduler = scheduler
        self.stop_request = threading.Event()
        self.report_interval = report_interval
        self.callbacks = callbacks
        self.max_workunit_verbosity = max_workunit_verbosity

    def run(self):
        while not self.stop_request.isSet():
            workunits = self.scheduler.poll_workunits(self.max_workunit_verbosity)
            for callback in self.callbacks:
                callback(
                    workunits=workunits["completed"],
                    started_workunits=workunits["started"],
                    finished=False,
                )
            self.stop_request.wait(timeout=self.report_interval)

    def join(self, timeout=None):
        self.stop_request.set()
        super(_InnerHandler, self).join(timeout)<commit_msg>Pass stream workunit context to all callbacks calls instead of just passing it in the last call. (#10340)<commit_after># Copyright 2019 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

import threading
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Callable, Iterable, Iterator, Optional, Sequence, Tuple

from pants.engine.fs import Digest
from pants.engine.internals.scheduler import SchedulerSession
from pants.util.logging import LogLevel


@dataclass(frozen=True)
class StreamingWorkunitContext:
    _scheduler: SchedulerSession

    def digests_to_bytes(self, digests: Sequence[Digest]) -> Tuple[bytes]:
        """Given a list of Digest objects, return a list of the bytes corresponding to each of those
        Digests in sequence."""
        return self._scheduler.digests_to_bytes(digests)

    def ensure_remote_has_recursive(self, digests: Sequence[Digest]) -> None:
        """Invoke the internal ensure_remote_has_recursive function, which ensures that a remote
        ByteStore, if it exists, has a copy of the files fingerprinted by each Digest."""
        return self._scheduler.ensure_remote_has_recursive(digests)


class StreamingWorkunitHandler:
    """StreamingWorkunitHandler's job is to periodically call each registered callback function with
    the following kwargs:

    workunits: Tuple[Dict[str, str], ...] - the workunit data itself
    finished: bool - this will be set to True when the last chunk of workunit data is reported to the callback
    """

    def __init__(
        self,
        scheduler: Any,
        callbacks: Iterable[Callable],
        report_interval_seconds: float,
        max_workunit_verbosity: LogLevel = LogLevel.DEBUG,
    ):
        self.scheduler = scheduler
        self.report_interval = report_interval_seconds
        self.callbacks = callbacks
        self._thread_runner: Optional[_InnerHandler] = None
        self._context = StreamingWorkunitContext(_scheduler=self.scheduler)
        # TODO(10092) The max verbosity should be a per-client setting, rather than a global setting.
        self.max_workunit_verbosity = max_workunit_verbosity

    def start(self) -> None:
        if self.callbacks:
            self._thread_runner = _InnerHandler(
                scheduler=self.scheduler,
                context=self._context,
                callbacks=self.callbacks,
                report_interval=self.report_interval,
                max_workunit_verbosity=self.max_workunit_verbosity,
            )
            self._thread_runner.start()

    def end(self) -> None:
        if self._thread_runner:
            self._thread_runner.join()

        # After stopping the thread, poll workunits one last time to make sure
        # we report any workunits that were added after the last time the thread polled.
        workunits = self.scheduler.poll_workunits(self.max_workunit_verbosity)
        for callback in self.callbacks:
            callback(
                workunits=workunits["completed"],
                started_workunits=workunits["started"],
                finished=True,
                context=self._context,
            )

    @contextmanager
    def session(self) -> Iterator[None]:
        try:
            self.start()
            yield
            self.end()
        except Exception as e:
            if self._thread_runner:
                self._thread_runner.join()
            raise e


class _InnerHandler(threading.Thread):
    def __init__(
        self,
        scheduler: Any,
        context: StreamingWorkunitContext,
        callbacks: Iterable[Callable],
        report_interval: float,
        max_workunit_verbosity: LogLevel,
    ):
        super().__init__(daemon=True)
        self.scheduler = scheduler
        self._context = context
        self.stop_request = threading.Event()
        self.report_interval = report_interval
        self.callbacks = callbacks
        self.max_workunit_verbosity = max_workunit_verbosity

    def run(self):
        while not self.stop_request.isSet():
            workunits = self.scheduler.poll_workunits(self.max_workunit_verbosity)
            for callback in self.callbacks:
                callback(
                    workunits=workunits["completed"],
                    started_workunits=workunits["started"],
                    finished=False,
                    context=self._context,
                )
            self.stop_request.wait(timeout=self.report_interval)

    def join(self, timeout=None):
        self.stop_request.set()
        super(_InnerHandler, self).join(timeout)
{'fraction_non_alphanumeric': 0.0730041428731385, 'fraction_numerical': 0.0030231776956667787, 'mean_word_length': 2.5813953488372094, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 13, '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': '7388551', 'n_tokens_mistral': 2419, 'n_tokens_neox': 2309, 'n_words': 644}