# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02_runners.ipynb (unless otherwise specified).


from __future__ import annotations


__all__ = ['ProcessMemory']

# Cell
#nbdev_comment from __future__ import annotations
from dataclasses import dataclass
from typing import Callable



# Cell
import logging
from tqdm.auto import tqdm

@dataclass
class ProcessMemory:
    queries: list[Callable]
    steps: list[Callable]
    filter: Callable
    progress_bar: bool = True

    def query(self):
        for query in tqdm(self.queries, desc='query', disable=not self.progress_bar):
            for record in query.query():
                yield record

    def prepare(self, records):
        return self.filter(tqdm(records, desc='filter', disable=not self.progress_bar))

    def fetch(self, records):
        for record in tqdm(records, desc='fetch', disable=not self.progress_bar):
            yield (record.content, record)

    def transform(self, content_records):
        for content, record in tqdm(content_records, desc='transform', disable=not self.progress_bar):
            error = False
            for step in self.steps:
                try:
                    content = step(content, record)
                except Exception as e:
                    error = True
                    logging.error('Error processing %s at step %s: %s' % (record, step.__name__, e))
                    break
            if not error:
                yield content

    def run(self):
        records = self.prepare(self.query())
        content_records = self.fetch(records)
        return self.transform(content_records)