<commit_before>import contextlib
import os
import tempfile

import datetime

from couchdbkit import ResourceConflict

from soil import DownloadBase

from couchexport.export import FormattedRow, get_writer
from couchexport.files import Temp
from couchexport.models import Format
from corehq.apps.export.esaccessors import (
    get_form_export_base_query,
    get_case_export_base_query,
)
from corehq.apps.export.models.new import (
    CaseExportInstance,
    FormExportInstance,
)


class ExportFile(object):
    # This is essentially coppied from couchexport.files.ExportFiles

    def __init__(self, path, format):
        self.file = Temp(path)
        self.format = format

    def __enter__(self):
        return self.file.payload

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.file.delete()


class _Writer(object):
    """
    An object that provides a friendlier interface to couchexport.ExportWriters.
    """
    def __init__(self, writer):
        # An instance of a couchexport.ExportWriter
        self.writer = writer
        self.format = writer.format
        self._path = None

    @contextlib.contextmanager
    def open(self, tables):
        """
        Open the _Writer for writing. This must be called before using _Writer.write().
        Note that this function returns a context manager!
        A _Writer can only be opened once.
        """

        # Create and open a temp file
        assert self._path is None
        fd, self._path = tempfile.mkstemp()
        with os.fdopen(fd, 'wb') as file:

            # open the ExportWriter
            headers = [(t, (t.get_headers(),)) for t in tables]
            table_titles = {t: t.label for t in tables}
            self.writer.open(headers, file, table_titles=table_titles)
            yield
            self.writer.close()

    def write(self, table, row):
        """
        Write the given row to the given table of the export.
        _Writer must be opened first.
        :param table: A TableConfiguration
        :param row: An ExportRow
        """
        return self.writer.write([(table, [FormattedRow(data=row.data)])])

    def get_preview(self):
        return self.writer.get_preview()

    @property
    def path(self):
        """
        The path to the file that this object writes to.
        """
        return self._path


def _get_writer(export_instances):
    """
    Return a new _Writer
    """
    format = Format.XLS_2007
    if len(export_instances) == 1:
        format = export_instances[0].export_format

    legacy_writer = get_writer(format)
    writer = _Writer(legacy_writer)
    return writer


def _get_tables(export_instances):
    """
    Return a list of tables for the given ExportInstances
    :param export_instances: A list of ExportInstances
    :return: a list of TableConfigurations
    """
    tables = []
    for export_instance in export_instances:
        tables.extend(export_instance.tables)
    return tables


def get_export_download(export_instances, filters, filename=None):
    from corehq.apps.export.tasks import populate_export_download_task

    download = DownloadBase()
    download.set_task(populate_export_download_task.delay(
        export_instances,
        filters,
        download.download_id,
        filename=filename
    ))
    return download


def get_export_file(export_instances, filters, progress_tracker=None):
    """
    Return an export file for the given ExportInstance and list of filters
    # TODO: Add a note about cleaning up the file?
    """

    writer = _get_writer(export_instances)
    with writer.open(_get_tables(export_instances)):
        for export_instance in export_instances:
            # TODO: Don't get the docs multiple times if you don't have to
            docs = _get_export_documents(export_instance, filters)
            _write_export_instance(writer, export_instance, docs, progress_tracker)

    return ExportFile(writer.path, writer.format)


def _get_export_documents(export_instance, filters):
    query = _get_base_query(export_instance)
    for filter in filters:
        query = query.filter(filter.to_es_filter())
    return query.scroll()


def _write_export_instance(writer, export_instance, documents, progress_tracker=None):
    """
    Write rows to the given open _Writer.<commit_msg>Fetch export docs in smaller chunks than the default 1 million<commit_after>import contextlib
import os
import tempfile

import datetime

from couchdbkit import ResourceConflict

from soil import DownloadBase

from couchexport.export import FormattedRow, get_writer
from couchexport.files import Temp
from couchexport.models import Format
from corehq.apps.export.esaccessors import (
    get_form_export_base_query,
    get_case_export_base_query,
)
from corehq.apps.export.models.new import (
    CaseExportInstance,
    FormExportInstance,
)


class ExportFile(object):
    # This is essentially coppied from couchexport.files.ExportFiles

    def __init__(self, path, format):
        self.file = Temp(path)
        self.format = format

    def __enter__(self):
        return self.file.payload

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.file.delete()


class _Writer(object):
    """
    An object that provides a friendlier interface to couchexport.ExportWriters.
    """
    def __init__(self, writer):
        # An instance of a couchexport.ExportWriter
        self.writer = writer
        self.format = writer.format
        self._path = None

    @contextlib.contextmanager
    def open(self, tables):
        """
        Open the _Writer for writing. This must be called before using _Writer.write().
        Note that this function returns a context manager!
        A _Writer can only be opened once.
        """

        # Create and open a temp file
        assert self._path is None
        fd, self._path = tempfile.mkstemp()
        with os.fdopen(fd, 'wb') as file:

            # open the ExportWriter
            headers = [(t, (t.get_headers(),)) for t in tables]
            table_titles = {t: t.label for t in tables}
            self.writer.open(headers, file, table_titles=table_titles)
            yield
            self.writer.close()

    def write(self, table, row):
        """
        Write the given row to the given table of the export.
        _Writer must be opened first.
        :param table: A TableConfiguration
        :param row: An ExportRow
        """
        return self.writer.write([(table, [FormattedRow(data=row.data)])])

    def get_preview(self):
        return self.writer.get_preview()

    @property
    def path(self):
        """
        The path to the file that this object writes to.
        """
        return self._path


def _get_writer(export_instances):
    """
    Return a new _Writer
    """
    format = Format.XLS_2007
    if len(export_instances) == 1:
        format = export_instances[0].export_format

    legacy_writer = get_writer(format)
    writer = _Writer(legacy_writer)
    return writer


def _get_tables(export_instances):
    """
    Return a list of tables for the given ExportInstances
    :param export_instances: A list of ExportInstances
    :return: a list of TableConfigurations
    """
    tables = []
    for export_instance in export_instances:
        tables.extend(export_instance.tables)
    return tables


def get_export_download(export_instances, filters, filename=None):
    from corehq.apps.export.tasks import populate_export_download_task

    download = DownloadBase()
    download.set_task(populate_export_download_task.delay(
        export_instances,
        filters,
        download.download_id,
        filename=filename
    ))
    return download


def get_export_file(export_instances, filters, progress_tracker=None):
    """
    Return an export file for the given ExportInstance and list of filters
    # TODO: Add a note about cleaning up the file?
    """

    writer = _get_writer(export_instances)
    with writer.open(_get_tables(export_instances)):
        for export_instance in export_instances:
            # TODO: Don't get the docs multiple times if you don't have to
            docs = _get_export_documents(export_instance, filters)
            _write_export_instance(writer, export_instance, docs, progress_tracker)

    return ExportFile(writer.path, writer.format)


def _get_export_documents(export_instance, filters):
    query = _get_base_query(export_instance)
    for filter in filters:
        query = query.filter(filter.to_es_filter())
    # size here limits each scroll request, not the total number of results
    return query.size(100).scroll()


def _write_export_instance(writer, export_instance, documents, progress_tracker=None):
    """
    Write rows to the given open _Writer.
{'fraction_non_alphanumeric': 0.07678265830005704, 'fraction_numerical': 0.0018254420992584142, 'mean_word_length': 2.9012016021361817, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6859559', 'n_tokens_mistral': 2532, 'n_tokens_neox': 2443, 'n_words': 735}