<commit_before>import itertools
from functools import partial
import math

from typing import Any, Callable, Iterable

from rechunker.types import CopySpec, StagedCopySpec, Executor


Thunk = Callable[[], None]


class PythonExecutor(Executor[Thunk]):
    """An execution engine based on Python loops.

    Supports copying between any arrays that implement ``__getitem__`` and
    ``__setitem__`` for tuples of ``slice`` objects.

    Execution plans for PythonExecutor are functions that accept no arguments.
    """

    def prepare_plan(self, specs: Iterable[StagedCopySpec]) -> Thunk:
        tasks = []
        for staged_copy_spec in specs:
            for copy_spec in staged_copy_spec.stages:
                tasks.append(partial(_direct_copy_array, copy_spec))
        return partial(_execute_all, tasks)

    def execute_plan(self, plan: Thunk):
        plan()


def _direct_copy_array(copy_spec: CopySpec) -> None:
    """Direct copy between zarr arrays."""
    source_array, target_array, chunks = copy_spec
    shape = source_array.shape
    ranges = [range(math.ceil(s / c)) for s, c in zip(shape, chunks)]
    for indices in itertools.product(*ranges):
        key = tuple(slice(c * i, c * (i + 1)) for i, c in zip(indices, chunks))
        target_array[key] = source_array[key]


def _execute_all(tasks: Iterable[Callable[[], Any]]) -> None:
    for task in tasks:
        task()<commit_msg>Remove 'thunk' jargon from PythonExecutor<commit_after>import itertools
from functools import partial
import math

from typing import Callable, Iterable

from rechunker.types import CopySpec, StagedCopySpec, Executor


# PythonExecutor represents delayed execution tasks as functions that require
# no arguments.
Task = Callable[[], None]


class PythonExecutor(Executor[Task]):
    """An execution engine based on Python loops.

    Supports copying between any arrays that implement ``__getitem__`` and
    ``__setitem__`` for tuples of ``slice`` objects.

    Execution plans for PythonExecutor are functions that accept no arguments.
    """

    def prepare_plan(self, specs: Iterable[StagedCopySpec]) -> Task:
        tasks = []
        for staged_copy_spec in specs:
            for copy_spec in staged_copy_spec.stages:
                tasks.append(partial(_direct_copy_array, copy_spec))
        return partial(_execute_all, tasks)

    def execute_plan(self, plan: Task):
        plan()


def _direct_copy_array(copy_spec: CopySpec) -> None:
    """Direct copy between zarr arrays."""
    source_array, target_array, chunks = copy_spec
    shape = source_array.shape
    ranges = [range(math.ceil(s / c)) for s, c in zip(shape, chunks)]
    for indices in itertools.product(*ranges):
        key = tuple(slice(c * i, c * (i + 1)) for i, c in zip(indices, chunks))
        target_array[key] = source_array[key]


def _execute_all(tasks: Iterable[Task]) -> None:
    for task in tasks:
        task()
{'fraction_non_alphanumeric': 0.09773662551440329, 'fraction_numerical': 0.0006858710562414266, 'mean_word_length': 3.5507020280811235, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, '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': '6412411', 'n_tokens_mistral': 870, 'n_tokens_neox': 842, 'n_words': 278}