<commit_before>from functools import partial

from typing import Callable, Iterable, Tuple

from rechunker.executors.util import chunk_keys, split_into_direct_copies
from rechunker.types import CopySpec, Executor, ReadableArray, WriteableArray

import pywren_ibm_cloud as pywren
from pywren_ibm_cloud.executor import FunctionExecutor

# PywrenExecutor represents delayed execution tasks as functions that require
# a FunctionExecutor.
Task = Callable[[FunctionExecutor], None]


class PywrenExecutor(Executor[Task]):
    """An execution engine based on Pywren.

    Supports zarr arrays as inputs. Outputs must be zarr arrays.

    Any Pywren FunctionExecutor can be passed to the constructor. By default
    a Pywren `local_executor` will be used

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

    def __init__(self, pywren_function_executor: FunctionExecutor = None):
        self.pywren_function_executor = pywren_function_executor

    def prepare_plan(self, specs: Iterable[CopySpec]) -> Task:
        tasks = []
        for spec in specs:
            for direct_spec in split_into_direct_copies(spec):
                tasks.append(partial(_direct_array_copy, *direct_spec))
        return partial(_execute_all, tasks)

    def execute_plan(self, plan: Task):
        if self.pywren_function_executor is None:
            # No Pywren function executor specified, so use a local one, and shutdown after use
            with pywren.local_executor() as pywren_function_executor:<commit_msg>Work around Pywren error if ~/.pywren_config is missing<commit_after>from functools import partial

from typing import Callable, Iterable, Tuple

from rechunker.executors.util import chunk_keys, split_into_direct_copies
from rechunker.types import CopySpec, Executor, ReadableArray, WriteableArray

import pywren_ibm_cloud as pywren
from pywren_ibm_cloud.executor import FunctionExecutor

# PywrenExecutor represents delayed execution tasks as functions that require
# a FunctionExecutor.
Task = Callable[[FunctionExecutor], None]


class PywrenExecutor(Executor[Task]):
    """An execution engine based on Pywren.

    Supports zarr arrays as inputs. Outputs must be zarr arrays.

    Any Pywren FunctionExecutor can be passed to the constructor. By default
    a Pywren `local_executor` will be used

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

    def __init__(self, pywren_function_executor: FunctionExecutor = None):
        self.pywren_function_executor = pywren_function_executor

    def prepare_plan(self, specs: Iterable[CopySpec]) -> Task:
        tasks = []
        for spec in specs:
            for direct_spec in split_into_direct_copies(spec):
                tasks.append(partial(_direct_array_copy, *direct_spec))
        return partial(_execute_all, tasks)

    def execute_plan(self, plan: Task):
        if self.pywren_function_executor is None:
            # No Pywren function executor specified, so use a local one, and shutdown after use
            with pywren.local_executor(
                # Minimal config needed to avoid Pywren error if ~/.pywren_config is missing
                config={"pywren": {"storage_bucket": "unused"}}
            ) as pywren_function_executor:
{'fraction_non_alphanumeric': 0.05975609756097561, 'fraction_numerical': 0.0, 'mean_word_length': 3.810850439882698, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 5, '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': '791217', 'n_tokens_mistral': 936, 'n_tokens_neox': 908, 'n_words': 302}