<commit_before>"""General utility functions."""

import datetime as dt
import logging as lg
import sys
import unicodedata
from pathlib import Path

from . import _version
from . import settings

def citation():
    """
    Print the OSMnx package's citation information.

    Boeing, G. 2017. OSMnx: New Methods for Acquiring, Constructing, Analyzing,
    and Visualizing Complex Street Networks. Computers, Environment and Urban
    Systems, 65(126-139). https://doi.org/10.1016/j.compenvurbsys.2017.05.004

    Returns
    -------
    None
    """
    cite = (
        "Citation:\n\n"
        "Boeing, G. 2017. OSMnx: New Methods for Acquiring, "
        "Constructing, Analyzing, and Visualizing Complex Street "
        "Networks. Computers, Environment and Urban Systems, 65(126-139). "
        "https://doi.org/10.1016/j.compenvurbsys.2017.05.004\n\n"
        "BibTeX entry for LaTeX users:\n\n"
        "@article{boeing_osmnx_2017,\n"
        "    title = {{OSMnx}: {New} {Methods} for {Acquiring}, "
        "{Constructing}, {Analyzing}, and {Visualizing} {Complex} "
        "{Street} {Networks}},\n"
        "    volume = {65},\n"
        "    doi = {10.1016/j.compenvurbsys.2017.05.004},\n"
        "    number = {126-139},\n"
        "    journal = {Computers, Environment and Urban Systems},\n"
        "    author = {Boeing, Geoff},\n"
        "    year = {2017}\n"
        "}"
    )

    print(cite)

def ts(style="datetime", template=None):
    """
    Get current timestamp as string.

    Parameters
    ----------
    style : string {"datetime", "date", "time"}
        format the timestamp with this built-in template
    template : string
        if not None, format the timestamp with this template instead of one of
        the built-in styles

    Returns
    -------
    ts : string
        the string timestamp
    """
    if template is None:
        if style == "datetime":
            template = "{:%Y-%m-%d %H:%M:%S}"
        elif style == "date":
            template = "{:%Y-%m-%d}"
        elif style == "time":
            template = "{:%H:%M:%S}"
        else:
            raise ValueError(f'unrecognized timestamp style "{style}"')

    ts = template.format(dt.datetime.now())
    return ts

def config(
    all_oneway=settings.all_oneway,
    bidirectional_network_types=settings.bidirectional_network_types,
    cache_folder=settings.cache_folder,
    cache_only_mode=settings.cache_only_mode,
    data_folder=settings.data_folder,
    default_accept_language=settings.default_accept_language,
    default_access=settings.default_access,
    default_crs=settings.default_crs,
    default_referer=settings.default_referer,
    default_user_agent=settings.default_user_agent,
    elevation_provider=settings.elevation_provider,
    imgs_folder=settings.imgs_folder,
    log_console=settings.log_console,
    log_file=settings.log_file,
    log_filename=settings.log_filename,
    log_level=settings.log_level,
    log_name=settings.log_name,
    logs_folder=settings.logs_folder,
    max_query_area_size=settings.max_query_area_size,
    memory=settings.memory,
    nominatim_endpoint=settings.nominatim_endpoint,
    nominatim_key=settings.nominatim_key,
    osm_xml_node_attrs=settings.osm_xml_node_attrs,
    osm_xml_node_tags=settings.osm_xml_node_tags,
    osm_xml_way_attrs=settings.osm_xml_way_attrs,
    osm_xml_way_tags=settings.osm_xml_way_tags,
    overpass_endpoint=settings.overpass_endpoint,
    overpass_settings=settings.overpass_settings,
    timeout=settings.timeout,
    use_cache=settings.use_cache,
    useful_tags_node=settings.useful_tags_node,
    useful_tags_way=settings.useful_tags_way,
):
    """
    Configure OSMnx by setting the default global settings' values.

    Any parameters not passed by the caller are (re-)set to their original
    default values.

    Parameters
    ----------
    all_oneway : bool
        Only use if specifically saving to .osm XML file with save_graph_xml
        function. if True, forces all ways to be loaded as oneway ways,
        preserving the original order of nodes stored in the OSM way XML.
        This also retains original OSM string values for oneway attribute
        values, rather than converting them to a True/False bool.
    bidirectional_network_types : list
        network types for which a fully bidirectional graph will be created
    cache_folder : string or pathlib.Path
        path to folder in which to save/load HTTP response cache
    data_folder : string or pathlib.Path
        path to folder in which to save/load graph files by default
    cache_only_mode : bool
        If True, download network data from Overpass then raise a
        CacheOnlyModeInterrupt error for user to catch. This prevents graph
        building from taking place and instead just saves OSM response data to
        cache. Useful for sequentially caching lots of raw data (as you can
        only query Overpass one request at a time) then using the cache to
        quickly build many graphs simultaneously with multiprocessing.
    default_accept_language : string
        HTTP header accept-language
    default_access : string
        default filter for OSM "access" key
    default_crs : string
        default coordinate reference system to set when creating graphs
    default_referer : string
        HTTP header referer
    default_user_agent : string
        HTTP header user-agent
    elevation_provider : string {"google", "airmap"}
        the API provider to use for adding node elevations
    imgs_folder : string or pathlib.Path
        path to folder in which to save plot images by default
    log_file : bool
        if True, save log output to a file in logs_folder
    log_filename : string
        name of the log file, without file extension
    log_console : bool
        if True, print log output to the console (terminal window)
    log_level : int
        one of Python's logger.level constants
    log_name : string
        name of the logger
    logs_folder : string or pathlib.Path
        path to folder in which to save log files
    max_query_area_size : int
        maximum area for any part of the geometry in meters: any polygon
        bigger than this will get divided up for multiple queries to API
        (default 50km x 50km)
    memory : int
        Overpass server memory allocation size for the query, in bytes. If
        None, server will use its default allocation size. Use with caution.
    nominatim_endpoint : string
        the API endpoint to use for nominatim queries
    nominatim_key : string
        your API key, if you are using an endpoint that requires one
    osm_xml_node_attrs : list
        node attributes for saving .osm XML files with save_graph_xml function
    osm_xml_node_tags : list
        node tags for saving .osm XML files with save_graph_xml function
    osm_xml_way_attrs : list
        edge attributes for saving .osm XML files with save_graph_xml function
    osm_xml_way_tags : list
        edge tags for for saving .osm XML files with save_graph_xml function
    overpass_endpoint : string
        the API endpoint to use for overpass queries
    overpass_settings : string
        Settings string for overpass queries. For example, to query historical
        OSM data as of a certain date:
        `'[out:json][timeout:90][date:"2019-10-28T19:20:00Z"]'`.
        Use with caution.
    timeout : int
        the timeout interval for the HTTP request and for API to use while
        running the query
    use_cache : bool
        if True, cache HTTP responses locally instead of calling API
        repeatedly for the same request
    useful_tags_node : list
        OSM "node" tags to add as graph node attributes, when present
    useful_tags_way : list
        OSM "way" tags to add as graph edge attributes, when present

    Returns
    -------
    None
    """
    # set each global setting to the argument value
    settings.all_oneway = all_oneway
    settings.bidirectional_network_types = bidirectional_network_types
    settings.cache_folder = cache_folder
    settings.cache_only_mode = cache_only_mode
    settings.data_folder = data_folder
    settings.default_accept_language = default_accept_language
    settings.default_access = default_access
    settings.default_crs = default_crs
    settings.default_referer = default_referer
    settings.default_user_agent = default_user_agent
    settings.elevation_provider = elevation_provider
    settings.imgs_folder = imgs_folder
    settings.log_console = log_console
    settings.log_file = log_file
    settings.log_filename = log_filename
    settings.log_level = log_level
    settings.log_name = log_name
    settings.logs_folder = logs_folder
    settings.max_query_area_size = max_query_area_size
    settings.memory = memory
    settings.nominatim_endpoint = nominatim_endpoint
    settings.nominatim_key = nominatim_key
    settings.osm_xml_node_attrs = osm_xml_node_attrs
    settings.osm_xml_node_tags = osm_xml_node_tags
    settings.osm_xml_way_attrs = osm_xml_way_attrs
    settings.osm_xml_way_tags = osm_xml_way_tags
    settings.overpass_endpoint = overpass_endpoint
    settings.overpass_settings = overpass_settings
    settings.timeout = timeout
    settings.use_cache = use_cache
    settings.useful_tags_node = useful_tags_node
    settings.useful_tags_way = useful_tags_way

    log(f"Configured OSMnx {_version.__version__}")
    log(f"HTTP response caching is {'on' if settings.use_cache else 'off'}")

def log(message, level=None, name=None, filename=None):
    """
    Write a message to the logger.

    This logs to file and/or prints to the console (terminal), depending on
    the current configuration of settings.log_file and settings.log_console.

    Parameters
    ----------
    message : string
        the message to log
    level : int
        one of Python's logger.level constants
    name : string
        name of the logger
    filename : string
        name of the log file, without file extension

    Returns
    -------
    None
    """
    if level is None:
        level = settings.log_level
    if name is None:
        name = settings.log_name
    if filename is None:
        filename = settings.log_filename

    # if logging to file is turned on
    if settings.log_file:
        # get the current logger (or create a new one, if none), then log
        # message at requested level
        logger = _get_logger(level=level, name=name, filename=filename)
        if level == lg.DEBUG:
            logger.debug(message)
        elif level == lg.INFO:
            logger.info(message)
        elif level == lg.WARNING:
            logger.warning(message)
        elif level == lg.ERROR:
            logger.error(message)

    # if logging to console, convert message to ascii and print to console
    if settings.log_console:
        # capture current stdout, then switch it to the console, print the
        # message, then switch back to what had been the stdout. this prevents
        # logging to notebook in jupyter: instead, it goes to terminal
        standard_out = sys.stdout
        sys.stdout = sys.__stdout__

        # prepend timestamp
        message = f"{ts()} {message}"

        # convert to ascii so it doesn't break windows terminals
        message = (
            unicodedata.normalize("NFKD", str(message)).encode("ascii", errors="replace").decode()
        )
        print(message)
        sys.stdout = standard_out

def _get_logger(level=None, name=None, filename=None):
    """
    Create a logger or return the current one if already instantiated.

    Parameters
    ----------
    level : int
        one of Python's logger.level constants
    name : string
        name of the logger
    filename : string
        name of the log file, without file extension

    Returns
    -------
    logger : logging.logger
    """
    if level is None:
        level = settings.log_level
    if name is None:
        name = settings.log_name
    if filename is None:
        filename = settings.log_filename

    logger = lg.getLogger(name)

    # if a logger with this name is not already set up
    if not getattr(logger, "handler_set", None):

        # get today's date and construct a log filename
        log_filename = Path(settings.logs_folder) / f'{filename}_{ts(style="date")}.log'

        # if the logs folder does not already exist, create it
        log_filename.parent.mkdir(parents=True, exist_ok=True)

        # create file handler and log formatter and set them up
        handler = lg.FileHandler(log_filename, encoding="utf-8")
        formatter = lg.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s")
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        logger.setLevel(level)
        logger.handler_set = True

    return logger
<commit_msg>Adjust formatting of inline code block.<commit_after>"""General utility functions."""

import datetime as dt
import logging as lg
import sys
import unicodedata
from pathlib import Path

from . import _version
from . import settings

def citation():
    """
    Print the OSMnx package's citation information.

    Boeing, G. 2017. OSMnx: New Methods for Acquiring, Constructing, Analyzing,
    and Visualizing Complex Street Networks. Computers, Environment and Urban
    Systems, 65(126-139). https://doi.org/10.1016/j.compenvurbsys.2017.05.004

    Returns
    -------
    None
    """
    cite = (
        "Citation:\n\n"
        "Boeing, G. 2017. OSMnx: New Methods for Acquiring, "
        "Constructing, Analyzing, and Visualizing Complex Street "
        "Networks. Computers, Environment and Urban Systems, 65(126-139). "
        "https://doi.org/10.1016/j.compenvurbsys.2017.05.004\n\n"
        "BibTeX entry for LaTeX users:\n\n"
        "@article{boeing_osmnx_2017,\n"
        "    title = {{OSMnx}: {New} {Methods} for {Acquiring}, "
        "{Constructing}, {Analyzing}, and {Visualizing} {Complex} "
        "{Street} {Networks}},\n"
        "    volume = {65},\n"
        "    doi = {10.1016/j.compenvurbsys.2017.05.004},\n"
        "    number = {126-139},\n"
        "    journal = {Computers, Environment and Urban Systems},\n"
        "    author = {Boeing, Geoff},\n"
        "    year = {2017}\n"
        "}"
    )

    print(cite)

def ts(style="datetime", template=None):
    """
    Get current timestamp as string.

    Parameters
    ----------
    style : string {"datetime", "date", "time"}
        format the timestamp with this built-in template
    template : string
        if not None, format the timestamp with this template instead of one of
        the built-in styles

    Returns
    -------
    ts : string
        the string timestamp
    """
    if template is None:
        if style == "datetime":
            template = "{:%Y-%m-%d %H:%M:%S}"
        elif style == "date":
            template = "{:%Y-%m-%d}"
        elif style == "time":
            template = "{:%H:%M:%S}"
        else:
            raise ValueError(f'unrecognized timestamp style "{style}"')

    ts = template.format(dt.datetime.now())
    return ts

def config(
    all_oneway=settings.all_oneway,
    bidirectional_network_types=settings.bidirectional_network_types,
    cache_folder=settings.cache_folder,
    cache_only_mode=settings.cache_only_mode,
    data_folder=settings.data_folder,
    default_accept_language=settings.default_accept_language,
    default_access=settings.default_access,
    default_crs=settings.default_crs,
    default_referer=settings.default_referer,
    default_user_agent=settings.default_user_agent,
    elevation_provider=settings.elevation_provider,
    imgs_folder=settings.imgs_folder,
    log_console=settings.log_console,
    log_file=settings.log_file,
    log_filename=settings.log_filename,
    log_level=settings.log_level,
    log_name=settings.log_name,
    logs_folder=settings.logs_folder,
    max_query_area_size=settings.max_query_area_size,
    memory=settings.memory,
    nominatim_endpoint=settings.nominatim_endpoint,
    nominatim_key=settings.nominatim_key,
    osm_xml_node_attrs=settings.osm_xml_node_attrs,
    osm_xml_node_tags=settings.osm_xml_node_tags,
    osm_xml_way_attrs=settings.osm_xml_way_attrs,
    osm_xml_way_tags=settings.osm_xml_way_tags,
    overpass_endpoint=settings.overpass_endpoint,
    overpass_settings=settings.overpass_settings,
    timeout=settings.timeout,
    use_cache=settings.use_cache,
    useful_tags_node=settings.useful_tags_node,
    useful_tags_way=settings.useful_tags_way,
):
    """
    Configure OSMnx by setting the default global settings' values.

    Any parameters not passed by the caller are (re-)set to their original
    default values.

    Parameters
    ----------
    all_oneway : bool
        Only use if specifically saving to .osm XML file with save_graph_xml
        function. if True, forces all ways to be loaded as oneway ways,
        preserving the original order of nodes stored in the OSM way XML.
        This also retains original OSM string values for oneway attribute
        values, rather than converting them to a True/False bool.
    bidirectional_network_types : list
        network types for which a fully bidirectional graph will be created
    cache_folder : string or pathlib.Path
        path to folder in which to save/load HTTP response cache
    data_folder : string or pathlib.Path
        path to folder in which to save/load graph files by default
    cache_only_mode : bool
        If True, download network data from Overpass then raise a
        CacheOnlyModeInterrupt error for user to catch. This prevents graph
        building from taking place and instead just saves OSM response data to
        cache. Useful for sequentially caching lots of raw data (as you can
        only query Overpass one request at a time) then using the cache to
        quickly build many graphs simultaneously with multiprocessing.
    default_accept_language : string
        HTTP header accept-language
    default_access : string
        default filter for OSM "access" key
    default_crs : string
        default coordinate reference system to set when creating graphs
    default_referer : string
        HTTP header referer
    default_user_agent : string
        HTTP header user-agent
    elevation_provider : string {"google", "airmap"}
        the API provider to use for adding node elevations
    imgs_folder : string or pathlib.Path
        path to folder in which to save plot images by default
    log_file : bool
        if True, save log output to a file in logs_folder
    log_filename : string
        name of the log file, without file extension
    log_console : bool
        if True, print log output to the console (terminal window)
    log_level : int
        one of Python's logger.level constants
    log_name : string
        name of the logger
    logs_folder : string or pathlib.Path
        path to folder in which to save log files
    max_query_area_size : int
        maximum area for any part of the geometry in meters: any polygon
        bigger than this will get divided up for multiple queries to API
        (default 50km x 50km)
    memory : int
        Overpass server memory allocation size for the query, in bytes. If
        None, server will use its default allocation size. Use with caution.
    nominatim_endpoint : string
        the API endpoint to use for nominatim queries
    nominatim_key : string
        your API key, if you are using an endpoint that requires one
    osm_xml_node_attrs : list
        node attributes for saving .osm XML files with save_graph_xml function
    osm_xml_node_tags : list
        node tags for saving .osm XML files with save_graph_xml function
    osm_xml_way_attrs : list
        edge attributes for saving .osm XML files with save_graph_xml function
    osm_xml_way_tags : list
        edge tags for for saving .osm XML files with save_graph_xml function
    overpass_endpoint : string
        the API endpoint to use for overpass queries
    overpass_settings : string
        Settings string for overpass queries. For example, to query historical
        OSM data as of a certain date:
        ``[out:json][timeout:90][date:"2019-10-28T19:20:00Z"]``.
        Use with caution.
    timeout : int
        the timeout interval for the HTTP request and for API to use while
        running the query
    use_cache : bool
        if True, cache HTTP responses locally instead of calling API
        repeatedly for the same request
    useful_tags_node : list
        OSM "node" tags to add as graph node attributes, when present
    useful_tags_way : list
        OSM "way" tags to add as graph edge attributes, when present

    Returns
    -------
    None
    """
    # set each global setting to the argument value
    settings.all_oneway = all_oneway
    settings.bidirectional_network_types = bidirectional_network_types
    settings.cache_folder = cache_folder
    settings.cache_only_mode = cache_only_mode
    settings.data_folder = data_folder
    settings.default_accept_language = default_accept_language
    settings.default_access = default_access
    settings.default_crs = default_crs
    settings.default_referer = default_referer
    settings.default_user_agent = default_user_agent
    settings.elevation_provider = elevation_provider
    settings.imgs_folder = imgs_folder
    settings.log_console = log_console
    settings.log_file = log_file
    settings.log_filename = log_filename
    settings.log_level = log_level
    settings.log_name = log_name
    settings.logs_folder = logs_folder
    settings.max_query_area_size = max_query_area_size
    settings.memory = memory
    settings.nominatim_endpoint = nominatim_endpoint
    settings.nominatim_key = nominatim_key
    settings.osm_xml_node_attrs = osm_xml_node_attrs
    settings.osm_xml_node_tags = osm_xml_node_tags
    settings.osm_xml_way_attrs = osm_xml_way_attrs
    settings.osm_xml_way_tags = osm_xml_way_tags
    settings.overpass_endpoint = overpass_endpoint
    settings.overpass_settings = overpass_settings
    settings.timeout = timeout
    settings.use_cache = use_cache
    settings.useful_tags_node = useful_tags_node
    settings.useful_tags_way = useful_tags_way

    log(f"Configured OSMnx {_version.__version__}")
    log(f"HTTP response caching is {'on' if settings.use_cache else 'off'}")

def log(message, level=None, name=None, filename=None):
    """
    Write a message to the logger.

    This logs to file and/or prints to the console (terminal), depending on
    the current configuration of settings.log_file and settings.log_console.

    Parameters
    ----------
    message : string
        the message to log
    level : int
        one of Python's logger.level constants
    name : string
        name of the logger
    filename : string
        name of the log file, without file extension

    Returns
    -------
    None
    """
    if level is None:
        level = settings.log_level
    if name is None:
        name = settings.log_name
    if filename is None:
        filename = settings.log_filename

    # if logging to file is turned on
    if settings.log_file:
        # get the current logger (or create a new one, if none), then log
        # message at requested level
        logger = _get_logger(level=level, name=name, filename=filename)
        if level == lg.DEBUG:
            logger.debug(message)
        elif level == lg.INFO:
            logger.info(message)
        elif level == lg.WARNING:
            logger.warning(message)
        elif level == lg.ERROR:
            logger.error(message)

    # if logging to console, convert message to ascii and print to console
    if settings.log_console:
        # capture current stdout, then switch it to the console, print the
        # message, then switch back to what had been the stdout. this prevents
        # logging to notebook in jupyter: instead, it goes to terminal
        standard_out = sys.stdout
        sys.stdout = sys.__stdout__

        # prepend timestamp
        message = f"{ts()} {message}"

        # convert to ascii so it doesn't break windows terminals
        message = (
            unicodedata.normalize("NFKD", str(message)).encode("ascii", errors="replace").decode()
        )
        print(message)
        sys.stdout = standard_out

def _get_logger(level=None, name=None, filename=None):
    """
    Create a logger or return the current one if already instantiated.

    Parameters
    ----------
    level : int
        one of Python's logger.level constants
    name : string
        name of the logger
    filename : string
        name of the log file, without file extension

    Returns
    -------
    logger : logging.logger
    """
    if level is None:
        level = settings.log_level
    if name is None:
        name = settings.log_name
    if filename is None:
        filename = settings.log_filename

    logger = lg.getLogger(name)

    # if a logger with this name is not already set up
    if not getattr(logger, "handler_set", None):

        # get today's date and construct a log filename
        log_filename = Path(settings.logs_folder) / f'{filename}_{ts(style="date")}.log'

        # if the logs folder does not already exist, create it
        log_filename.parent.mkdir(parents=True, exist_ok=True)

        # create file handler and log formatter and set them up
        handler = lg.FileHandler(log_filename, encoding="utf-8")
        formatter = lg.Formatter("%(asctime)s %(levelname)s %(name)s %(message)s")
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        logger.setLevel(level)
        logger.handler_set = True

    return logger
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
import contextlib
import time
import sys
from collections import Counter

import datetime

from couchdbkit import ResourceConflict

from dimagi.utils.logging import notify_exception
from soil import DownloadBase

from couchexport.export import FormattedRow, get_writer
from couchexport.models import Format
from corehq.elastic import iter_es_docs_from_query
from corehq.toggles import PAGINATED_EXPORTS
from corehq.util.files import safe_filename, TransientTempfile
from corehq.util.datadog.gauges import datadog_histogram, datadog_track_errors
from corehq.apps.export.esaccessors import (
    get_form_export_base_query,
    get_case_export_base_query,
    get_sms_export_base_query,
)
from corehq.apps.export.models.new import (
    CaseExportInstance,
    FormExportInstance,
    SMSExportInstance,
)
from corehq.apps.export.const import MAX_EXPORTABLE_ROWS
from corehq.apps.export.dbaccessors import get_properly_wrapped_export_instance
import six
from io import open

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

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

    def __enter__(self):
        self.file = open(self.path, 'rb')
        return self.file

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

class _ExportWriter(object):
    """
    An object that provides a friendlier interface to couchexport.ExportWriters.
    """

    def __init__(self, writer, temp_path):
        # An instance of a couchexport.ExportWriter
        self.writer = writer
        self.format = writer.format
        self.path = temp_path

    @contextlib.contextmanager
    def open(self, export_instances):
        """
        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.
        """

        if len(export_instances) == 1:
            name = export_instances[0].name or ''
        else:
            name = ''
        name = safe_filename(name)
        all_sheet_names = Counter(
            table.label for instance in export_instances for table in instance.selected_tables
        )

        with open(self.path, 'wb') as file:

            # open the ExportWriter
            headers = []
            table_titles = {}
            for instance_index, instance in enumerate(export_instances):
                headers += [
                    (t, (t.get_headers(split_columns=instance.split_multiselects),))
                    for t in instance.selected_tables
                ]
                for table_index, table in enumerate(instance.selected_tables):
                    sheet_name = table.label or "Sheet{}".format(table_index + 1)
                    # If it's a bulk export and the sheet has the same name as another sheet,
                    # Prefix the sheet name with the export name
                    if len(export_instances) > 1 and all_sheet_names[sheet_name] > 1:
                        sheet_name = "{}-{}".format(
                            instance.name or "Export{}".format(instance_index + 1),
                            sheet_name
                        )
                    table_titles[table] = sheet_name
            self.writer.open(headers, file, table_titles=table_titles, archive_basepath=name)
            try:
                yield
            finally:
                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()

class _PaginatedExportWriter(object):

    def __init__(self, writer, temp_path):
        self.format = writer.format
        self.path = temp_path
        self.pages = Counter()
        self.rows_written = Counter()
        # An instance of a couchexport.ExportWriter
        self.writer = writer
        self.file_handle = None

    @contextlib.contextmanager
    def open(self, export_instances):
        """
        Open the _PaginatedWriter for writing. This must be called before using _PaginatedWriter.write().
        A _PaginatedWriter can only be opened once.
        """

        self.name = self._get_name(export_instances)
        self.headers = self._get_headers(export_instances)
        self.table_names = self._get_table_names(export_instances)

        with open(self.path, 'wb') as file_handle:
            self.writer.open(
                six.iteritems(self._get_paginated_headers()),
                file_handle,
                table_titles=self._get_paginated_table_titles(),
                archive_basepath=self.name
            )
            try:
                yield
            finally:
                self.writer.close()

    def _get_name(self, export_instances):
        if len(export_instances) == 1:
            name = export_instances[0].name or 'Export'
        else:
            name = 'BulkExport'
        return safe_filename(name)

    def _get_headers(self, export_instances):
        '''
        Returns a dictionary that maps all TableConfigurations in the list of ExportInstances to an
        array of headers
        {
            TableConfiguration(): ['Column1', 'Column2']
            ...
        }

        :export_instances: - A list of ExportInstances
        '''
        headers = {}
        for instance_index, instance in enumerate(export_instances):
            for table in instance.selected_tables:
                headers[table] = table.get_headers(
                    split_columns=instance.split_multiselects
                )

        return headers

    def _get_table_names(self, export_instances):
        '''
        Returns a dictionary that maps all TableConfigurations in the list of ExportInstances to a
        table name. If the table name is a duplicate from another ExportInstance, it will automatically
        prefix the table name with the export name to prevent duplicate table names.

        {
            TableConfiguration(): 'Table Name'
        }
        '''
        table_names = {}
        for instance_index, instance in enumerate(export_instances):
            for table_index, table in enumerate(instance.selected_tables):
                table_name = table.label or "Sheet{}".format(table_index + 1)
                if len(export_instances) > 1:
                    table_name = "{}-{}".format(
                        instance.name or "Export{}".format(instance_index + 1),
                        table_name
                    )
                table_names[table] = table_name
        return table_names

    def _paged_table_index(self, table):
        '''
        Returns an index of the paged table. Relies on the internal state of self.pages
        ((PathNode(), PathNode()), 3)
        '''
        return (tuple(table.path), self.pages[table])

    def _get_paginated_headers(self):
        '''
        Maps the headers in self.headers to
        1) Match the data structure that couchexport.writers.ExportWriter expects
        2) Map the key in the headers dictionary to the paged table index

        {
            TableConfiguration(): ['Column1', 'Column2']
        }

        Becomes

        {
            (<table.path>, <page>): (['Column1', 'Column2'],)
        }
        '''
        return {self._paged_table_index(table): (headers,) for table, headers in six.iteritems(self.headers)}

    def _get_paginated_table_titles(self):
        '''
        Maps the table titles to tables titles with their page count

        {
            TableConfiguration(): 'Table Name'
        }

        Becomes
        {
            (<table.path>, <page>: 'Table Name_<page>'
        }
        '''
        paginated_table_titles = {}
        for table, table_name in six.iteritems(self.table_names):
            paginated_table_titles[self._paged_table_index(table)] = '{}_{}'.format(
                table_name, format(self.pages[table], '03')
            )
        return paginated_table_titles

    def write(self, table, row):
        """
        Write the given row to the given table of the export.
        Will automatically open a new table and write to that if it
        has exceeded the number of rows written in the first table.
        :param table: A TableConfiguration
        :param row: An ExportRow
        """
        if self.rows_written[table] >= MAX_EXPORTABLE_ROWS * (self.pages[table] + 1):
            self.pages[table] += 1
            self.writer.add_table(
                self._paged_table_index(table),
                self._get_paginated_headers()[self._paged_table_index(table)][0],
                table_title=self._get_paginated_table_titles()[self._paged_table_index(table)],
            )

        self.writer.write([(self._paged_table_index(table), [FormattedRow(data=row.data)])])
        self.rows_written[table] += 1

def get_export_writer(export_instances, temp_path, allow_pagination=True):
    """
    Return a new _Writer
    """
    format = Format.XLS_2007
    if len(export_instances) == 1:
        format = export_instances[0].export_format

    legacy_writer = get_writer(format)
    if allow_pagination and PAGINATED_EXPORTS.enabled(export_instances[0].domain):
        writer = _PaginatedExportWriter(legacy_writer, temp_path)
    else:
        writer = _ExportWriter(legacy_writer, temp_path)
    return writer

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, temp_path, progress_tracker=None):
    """
    Return an export file for the given ExportInstance and list of filters
    """
    writer = get_export_writer(export_instances, temp_path)
    with writer.open(export_instances):
        for export_instance in export_instances:
            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):
    # Pull doc ids from elasticsearch and stream to disk
    query = _get_export_query(export_instance, filters)
    return iter_es_docs_from_query(query)

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

def get_export_size(export_instance, filters):
    return _get_export_query(export_instance, filters).count()

def write_export_instance(writer, export_instance, documents, progress_tracker=None):
    """
    Write rows to the given open _Writer.
    Rows will be written to each table in the export instance for each of
    the given documents.
    :param writer: An open _Writer
    :param export_instance: An ExportInstance
    :param documents: An iterable yielding documents
    :param progress_tracker: A task for soil to track progress against
    :return: None
    """
    if progress_tracker:
        DownloadBase.set_progress(progress_tracker, 0, documents.count)

    start = _time_in_milliseconds()
    total_bytes = 0
    total_rows = 0
    compute_total = 0
    write_total = 0

    for row_number, doc in enumerate(documents):
        total_bytes += sys.getsizeof(doc)
        for table in export_instance.selected_tables:
            compute_start = _time_in_milliseconds()
            try:
                rows = table.get_rows(
                    doc,
                    row_number,
                    split_columns=export_instance.split_multiselects,
                    transform_dates=export_instance.transform_dates,
                )
            except Exception as e:
                notify_exception(None, "Error exporting doc", details={
                    'domain': export_instance.domain,
                    'export_instance_id': export_instance.get_id,
                    'export_table': table.label,
                    'doc_id': doc.get('_id'),
                })
                e.sentry_capture = False
                raise
            compute_total += _time_in_milliseconds() - compute_start

            write_start = _time_in_milliseconds()
            for row in rows:
                # It might be bad to write one row at a time when you can do more (from a performance perspective)
                # Regardless, we should handle the batching of rows in the _Writer class, not here.
                writer.write(table, row)
            write_total += _time_in_milliseconds() - write_start

            total_rows += len(rows)

        if progress_tracker:
            DownloadBase.set_progress(progress_tracker, row_number + 1, documents.count)

    end = _time_in_milliseconds()
    tags = ['format:{}'.format(writer.format), 'export_id:{}'.format(export_instance.get_id)]
    _record_datadog_export_write_rows(write_total, total_bytes, total_rows, tags)
    _record_datadog_export_compute_rows(compute_total, total_bytes, total_rows, tags)
    _record_datadog_export_duration(end - start, total_bytes, total_rows, tags)

def _time_in_milliseconds():
    return int(time.time() * 1000)

def _record_datadog_export_compute_rows(duration, doc_bytes, n_rows, tags):
    __record_datadog_export(duration, doc_bytes, n_rows, 'commcare.compute_export_rows_duration', tags)

def _record_datadog_export_write_rows(duration, doc_bytes, n_rows, tags):
    __record_datadog_export(duration, doc_bytes, n_rows, 'commcare.write_export_rows_duration', tags)

def _record_datadog_export_duration(duration, doc_bytes, n_rows, tags):
    __record_datadog_export(duration, doc_bytes, n_rows, 'commcare.export_duration', tags)

def __record_datadog_export(duration, doc_bytes, n_rows, metric, tags):
    datadog_histogram(metric, duration, tags=tags)
    if doc_bytes:
        datadog_histogram(
            '{}_normalized_by_size'.format(metric),
            duration // doc_bytes,
            tags=tags,
        )
    if n_rows:
        datadog_histogram(
            '{}_normalized_by_rows'.format(metric),
            duration // n_rows,
            tags=tags,
        )

def _get_base_query(export_instance):
    """
    Return an ESQuery object for the given export instance.
    Includes filters for domain, doc_type, and xmlns/case_type.
    """
    if isinstance(export_instance, FormExportInstance):
        return get_form_export_base_query(
            export_instance.domain,
            export_instance.app_id,
            export_instance.xmlns,
            export_instance.include_errors
        )
    if isinstance(export_instance, CaseExportInstance):
        return get_case_export_base_query(
            export_instance.domain, export_instance.case_type
        )
    if isinstance(export_instance, SMSExportInstance):
        return get_sms_export_base_query(export_instance.domain)
    else:
        raise Exception(
            "Unknown base query for export instance type {}".format(type(export_instance))
        )

@datadog_track_errors('rebuild_export')
def rebuild_export(export_instance, progress_tracker):
    """
    Rebuild the given daily saved ExportInstance
    """
    filters = export_instance.get_filters()
    with TransientTempfile() as temp_path:
        export_file = get_export_file([export_instance], filters or [], temp_path, progress_tracker)
        with export_file as payload:
            save_export_payload(export_instance, payload)

def save_export_payload(export, payload):
    """
    Save the contents of an export file to disk for later retrieval.
    """
    if export.last_accessed is None:
        export.last_accessed = datetime.datetime.utcnow()
    export.last_updated = datetime.datetime.utcnow()

    try:
        with export.atomic_blobs():
            export.set_payload(payload)
    except ResourceConflict:
        # task was executed concurrently, so let first to finish win and abort the rest
        pass
<commit_msg>lint[skip ci]<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
import contextlib
import time
import sys
from collections import Counter

import datetime

from couchdbkit import ResourceConflict

from dimagi.utils.logging import notify_exception
from soil import DownloadBase

from couchexport.export import FormattedRow, get_writer
from couchexport.models import Format
from corehq.elastic import iter_es_docs_from_query
from corehq.toggles import PAGINATED_EXPORTS
from corehq.util.files import safe_filename, TransientTempfile
from corehq.util.datadog.gauges import datadog_histogram, datadog_track_errors
from corehq.apps.export.esaccessors import (
    get_form_export_base_query,
    get_case_export_base_query,
    get_sms_export_base_query,
)
from corehq.apps.export.models.new import (
    CaseExportInstance,
    FormExportInstance,
    SMSExportInstance,
)
from corehq.apps.export.const import MAX_EXPORTABLE_ROWS
import six
from io import open

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

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

    def __enter__(self):
        self.file = open(self.path, 'rb')
        return self.file

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

class _ExportWriter(object):
    """
    An object that provides a friendlier interface to couchexport.ExportWriters.
    """

    def __init__(self, writer, temp_path):
        # An instance of a couchexport.ExportWriter
        self.writer = writer
        self.format = writer.format
        self.path = temp_path

    @contextlib.contextmanager
    def open(self, export_instances):
        """
        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.
        """

        if len(export_instances) == 1:
            name = export_instances[0].name or ''
        else:
            name = ''
        name = safe_filename(name)
        all_sheet_names = Counter(
            table.label for instance in export_instances for table in instance.selected_tables
        )

        with open(self.path, 'wb') as file:

            # open the ExportWriter
            headers = []
            table_titles = {}
            for instance_index, instance in enumerate(export_instances):
                headers += [
                    (t, (t.get_headers(split_columns=instance.split_multiselects),))
                    for t in instance.selected_tables
                ]
                for table_index, table in enumerate(instance.selected_tables):
                    sheet_name = table.label or "Sheet{}".format(table_index + 1)
                    # If it's a bulk export and the sheet has the same name as another sheet,
                    # Prefix the sheet name with the export name
                    if len(export_instances) > 1 and all_sheet_names[sheet_name] > 1:
                        sheet_name = "{}-{}".format(
                            instance.name or "Export{}".format(instance_index + 1),
                            sheet_name
                        )
                    table_titles[table] = sheet_name
            self.writer.open(headers, file, table_titles=table_titles, archive_basepath=name)
            try:
                yield
            finally:
                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()

class _PaginatedExportWriter(object):

    def __init__(self, writer, temp_path):
        self.format = writer.format
        self.path = temp_path
        self.pages = Counter()
        self.rows_written = Counter()
        # An instance of a couchexport.ExportWriter
        self.writer = writer
        self.file_handle = None

    @contextlib.contextmanager
    def open(self, export_instances):
        """
        Open the _PaginatedWriter for writing. This must be called before using _PaginatedWriter.write().
        A _PaginatedWriter can only be opened once.
        """

        self.name = self._get_name(export_instances)
        self.headers = self._get_headers(export_instances)
        self.table_names = self._get_table_names(export_instances)

        with open(self.path, 'wb') as file_handle:
            self.writer.open(
                six.iteritems(self._get_paginated_headers()),
                file_handle,
                table_titles=self._get_paginated_table_titles(),
                archive_basepath=self.name
            )
            try:
                yield
            finally:
                self.writer.close()

    def _get_name(self, export_instances):
        if len(export_instances) == 1:
            name = export_instances[0].name or 'Export'
        else:
            name = 'BulkExport'
        return safe_filename(name)

    def _get_headers(self, export_instances):
        '''
        Returns a dictionary that maps all TableConfigurations in the list of ExportInstances to an
        array of headers
        {
            TableConfiguration(): ['Column1', 'Column2']
            ...
        }

        :export_instances: - A list of ExportInstances
        '''
        headers = {}
        for instance_index, instance in enumerate(export_instances):
            for table in instance.selected_tables:
                headers[table] = table.get_headers(
                    split_columns=instance.split_multiselects
                )

        return headers

    def _get_table_names(self, export_instances):
        '''
        Returns a dictionary that maps all TableConfigurations in the list of ExportInstances to a
        table name. If the table name is a duplicate from another ExportInstance, it will automatically
        prefix the table name with the export name to prevent duplicate table names.

        {
            TableConfiguration(): 'Table Name'
        }
        '''
        table_names = {}
        for instance_index, instance in enumerate(export_instances):
            for table_index, table in enumerate(instance.selected_tables):
                table_name = table.label or "Sheet{}".format(table_index + 1)
                if len(export_instances) > 1:
                    table_name = "{}-{}".format(
                        instance.name or "Export{}".format(instance_index + 1),
                        table_name
                    )
                table_names[table] = table_name
        return table_names

    def _paged_table_index(self, table):
        '''
        Returns an index of the paged table. Relies on the internal state of self.pages
        ((PathNode(), PathNode()), 3)
        '''
        return (tuple(table.path), self.pages[table])

    def _get_paginated_headers(self):
        '''
        Maps the headers in self.headers to
        1) Match the data structure that couchexport.writers.ExportWriter expects
        2) Map the key in the headers dictionary to the paged table index

        {
            TableConfiguration(): ['Column1', 'Column2']
        }

        Becomes

        {
            (<table.path>, <page>): (['Column1', 'Column2'],)
        }
        '''
        return {self._paged_table_index(table): (headers,) for table, headers in six.iteritems(self.headers)}

    def _get_paginated_table_titles(self):
        '''
        Maps the table titles to tables titles with their page count

        {
            TableConfiguration(): 'Table Name'
        }

        Becomes
        {
            (<table.path>, <page>: 'Table Name_<page>'
        }
        '''
        paginated_table_titles = {}
        for table, table_name in six.iteritems(self.table_names):
            paginated_table_titles[self._paged_table_index(table)] = '{}_{}'.format(
                table_name, format(self.pages[table], '03')
            )
        return paginated_table_titles

    def write(self, table, row):
        """
        Write the given row to the given table of the export.
        Will automatically open a new table and write to that if it
        has exceeded the number of rows written in the first table.
        :param table: A TableConfiguration
        :param row: An ExportRow
        """
        if self.rows_written[table] >= MAX_EXPORTABLE_ROWS * (self.pages[table] + 1):
            self.pages[table] += 1
            self.writer.add_table(
                self._paged_table_index(table),
                self._get_paginated_headers()[self._paged_table_index(table)][0],
                table_title=self._get_paginated_table_titles()[self._paged_table_index(table)],
            )

        self.writer.write([(self._paged_table_index(table), [FormattedRow(data=row.data)])])
        self.rows_written[table] += 1

def get_export_writer(export_instances, temp_path, allow_pagination=True):
    """
    Return a new _Writer
    """
    format = Format.XLS_2007
    if len(export_instances) == 1:
        format = export_instances[0].export_format

    legacy_writer = get_writer(format)
    if allow_pagination and PAGINATED_EXPORTS.enabled(export_instances[0].domain):
        writer = _PaginatedExportWriter(legacy_writer, temp_path)
    else:
        writer = _ExportWriter(legacy_writer, temp_path)
    return writer

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, temp_path, progress_tracker=None):
    """
    Return an export file for the given ExportInstance and list of filters
    """
    writer = get_export_writer(export_instances, temp_path)
    with writer.open(export_instances):
        for export_instance in export_instances:
            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):
    # Pull doc ids from elasticsearch and stream to disk
    query = _get_export_query(export_instance, filters)
    return iter_es_docs_from_query(query)

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

def get_export_size(export_instance, filters):
    return _get_export_query(export_instance, filters).count()

def write_export_instance(writer, export_instance, documents, progress_tracker=None):
    """
    Write rows to the given open _Writer.
    Rows will be written to each table in the export instance for each of
    the given documents.
    :param writer: An open _Writer
    :param export_instance: An ExportInstance
    :param documents: An iterable yielding documents
    :param progress_tracker: A task for soil to track progress against
    :return: None
    """
    if progress_tracker:
        DownloadBase.set_progress(progress_tracker, 0, documents.count)

    start = _time_in_milliseconds()
    total_bytes = 0
    total_rows = 0
    compute_total = 0
    write_total = 0

    for row_number, doc in enumerate(documents):
        total_bytes += sys.getsizeof(doc)
        for table in export_instance.selected_tables:
            compute_start = _time_in_milliseconds()
            try:
                rows = table.get_rows(
                    doc,
                    row_number,
                    split_columns=export_instance.split_multiselects,
                    transform_dates=export_instance.transform_dates,
                )
            except Exception as e:
                notify_exception(None, "Error exporting doc", details={
                    'domain': export_instance.domain,
                    'export_instance_id': export_instance.get_id,
                    'export_table': table.label,
                    'doc_id': doc.get('_id'),
                })
                e.sentry_capture = False
                raise
            compute_total += _time_in_milliseconds() - compute_start

            write_start = _time_in_milliseconds()
            for row in rows:
                # It might be bad to write one row at a time when you can do more (from a performance perspective)
                # Regardless, we should handle the batching of rows in the _Writer class, not here.
                writer.write(table, row)
            write_total += _time_in_milliseconds() - write_start

            total_rows += len(rows)

        if progress_tracker:
            DownloadBase.set_progress(progress_tracker, row_number + 1, documents.count)

    end = _time_in_milliseconds()
    tags = ['format:{}'.format(writer.format), 'export_id:{}'.format(export_instance.get_id)]
    _record_datadog_export_write_rows(write_total, total_bytes, total_rows, tags)
    _record_datadog_export_compute_rows(compute_total, total_bytes, total_rows, tags)
    _record_datadog_export_duration(end - start, total_bytes, total_rows, tags)

def _time_in_milliseconds():
    return int(time.time() * 1000)

def _record_datadog_export_compute_rows(duration, doc_bytes, n_rows, tags):
    __record_datadog_export(duration, doc_bytes, n_rows, 'commcare.compute_export_rows_duration', tags)

def _record_datadog_export_write_rows(duration, doc_bytes, n_rows, tags):
    __record_datadog_export(duration, doc_bytes, n_rows, 'commcare.write_export_rows_duration', tags)

def _record_datadog_export_duration(duration, doc_bytes, n_rows, tags):
    __record_datadog_export(duration, doc_bytes, n_rows, 'commcare.export_duration', tags)

def __record_datadog_export(duration, doc_bytes, n_rows, metric, tags):
    datadog_histogram(metric, duration, tags=tags)
    if doc_bytes:
        datadog_histogram(
            '{}_normalized_by_size'.format(metric),
            duration // doc_bytes,
            tags=tags,
        )
    if n_rows:
        datadog_histogram(
            '{}_normalized_by_rows'.format(metric),
            duration // n_rows,
            tags=tags,
        )

def _get_base_query(export_instance):
    """
    Return an ESQuery object for the given export instance.
    Includes filters for domain, doc_type, and xmlns/case_type.
    """
    if isinstance(export_instance, FormExportInstance):
        return get_form_export_base_query(
            export_instance.domain,
            export_instance.app_id,
            export_instance.xmlns,
            export_instance.include_errors
        )
    if isinstance(export_instance, CaseExportInstance):
        return get_case_export_base_query(
            export_instance.domain, export_instance.case_type
        )
    if isinstance(export_instance, SMSExportInstance):
        return get_sms_export_base_query(export_instance.domain)
    else:
        raise Exception(
            "Unknown base query for export instance type {}".format(type(export_instance))
        )

@datadog_track_errors('rebuild_export')
def rebuild_export(export_instance, progress_tracker):
    """
    Rebuild the given daily saved ExportInstance
    """
    filters = export_instance.get_filters()
    with TransientTempfile() as temp_path:
        export_file = get_export_file([export_instance], filters or [], temp_path, progress_tracker)
        with export_file as payload:
            save_export_payload(export_instance, payload)

def save_export_payload(export, payload):
    """
    Save the contents of an export file to disk for later retrieval.
    """
    if export.last_accessed is None:
        export.last_accessed = datetime.datetime.utcnow()
    export.last_updated = datetime.datetime.utcnow()

    try:
        with export.atomic_blobs():
            export.set_payload(payload)
    except ResourceConflict:
        # task was executed concurrently, so let first to finish win and abort the rest
        pass
<|endoftext|>


<commit_before>from django import forms
from django.contrib import auth

from .models import AccountRecovery

class ConfirmForm(forms.Form):
    pass

class UserCreationForm(auth.forms.UserCreationForm):
    def clean_username(self):
        # Since User.username is unique, this check is redundant,
        # but it sets a nicer error message than the ORM. See #13147.
        username = self.cleaned_data["username"]
        try:
            self.Meta.model._default_manager.get(username=username)
        except self.Meta.model.DoesNotExist:
            return username
        raise forms.ValidationError(
            self.error_messages['duplicate_username'],
            code='duplicate_username',
        )

    class Meta:
        model = auth.get_user_model()
        fields = ()

class UserAuthenticationForm(auth.forms.AuthenticationForm):
    pass
<commit_msg>Username is selected as a field<commit_after>from django import forms
from django.contrib import auth

from .models import AccountRecovery

class ConfirmForm(forms.Form):
    pass

class UserCreationForm(auth.forms.UserCreationForm):
    def clean_username(self):
        # Since User.username is unique, this check is redundant,
        # but it sets a nicer error message than the ORM. See #13147.
        username = self.cleaned_data["username"]
        try:
            self.Meta.model._default_manager.get(username=username)
        except self.Meta.model.DoesNotExist:
            return username
        raise forms.ValidationError(
            self.error_messages['duplicate_username'],
            code='duplicate_username',
        )

    class Meta:
        model = auth.get_user_model()
        fields = ('username', )

class UserAuthenticationForm(auth.forms.AuthenticationForm):
    pass
<|endoftext|>


<commit_before>import json
import os
from datetime import datetime
from mock import patch
from django.test import SimpleTestCase
from corehq.apps.app_manager.models import Application
from corehq.apps.userreports.app_manager import get_case_data_sources, get_form_data_sources
from corehq.apps.userreports.reports.builder import DEFAULT_CASE_PROPERTY_DATATYPES

class AppManagerDataSourceConfigTest(SimpleTestCase):

    def setUp(self):
        self.is_usercase_in_use_patch = patch('corehq.apps.app_manager.models.is_usercase_in_use')
        is_usercase_in_use_mock = self.is_usercase_in_use_patch.start()
        is_usercase_in_use_mock.return_value = False

    def tearDown(self):
        self.is_usercase_in_use_patch.stop()

    def get_json(self, name):
        with open(os.path.join(os.path.dirname(__file__), 'data', 'app_manager', name)) as f:
            return json.loads(f.read())

    @patch('corehq.apps.userreports.specs.datetime')
    def test_simple_case_management(self, datetime_mock):
        fake_time_now = datetime(2015, 4, 24, 12, 30, 8, 24886)
        datetime_mock.utcnow.return_value = fake_time_now

        app = Application.wrap(self.get_json('simple_app.json'))
        self.assertEqual('userreports_test', app.domain)
        data_sources = get_case_data_sources(app)
        self.assertEqual(1, len(data_sources))
        data_source = data_sources['ticket']
        self.assertEqual(app.domain, data_source.domain)
        self.assertEqual('CommCareCase', data_source.referenced_doc_type)
        self.assertEqual('ticket', data_source.table_id)
        self.assertEqual('ticket', data_source.display_name)

        # test the filter
        self.assertTrue(data_source.filter(
            {'doc_type': 'CommCareCase', 'domain': app.domain, 'type': 'ticket'}))
        self.assertFalse(data_source.filter(
            {'doc_type': 'CommCareCase', 'domain': 'wrong domain', 'type': 'ticket'}))
        self.assertFalse(data_source.filter(
            {'doc_type': 'NotCommCareCase', 'domain': app.domain, 'type': 'ticket'}))
        self.assertFalse(data_source.filter(
            {'doc_type': 'CommCareCase', 'domain': app.domain, 'type': 'not-ticket'}))

        # check the indicators
        expected_columns = set(
            ["doc_id", "modified_on", "user_id", "opened_on",
             "owner_id", "name", "category", "priority", "starred", "estimate", "inserted_at"]
        )
        self.assertEqual(expected_columns, set(col_back.id for col_back in data_source.get_columns()))

        modified_on = datetime(2014, 11, 12, 15, 37, 49)
        opened_on = datetime(2014, 11, 11, 23, 34, 34, 25)
        sample_doc = dict(
            _id='some-doc-id',
            doc_type="CommCareCase",
            modified_on=modified_on.isoformat() + 'Z',
            user_id="23407238074",
            opened_on=opened_on.isoformat() + 'Z',
            owner_id="0923409230948",
            name="priority ticket",
            domain=app.domain,
            type='ticket',
            category='bug',
            priority='4',
            starred='yes',
            estimate='2',
        )

        def _get_column_property(column):
            return column.id if column.id != 'doc_id' else '_id'

        default_case_property_datatypes = DEFAULT_CASE_PROPERTY_DATATYPES
        [row] = data_source.get_all_values(sample_doc)
        for result in row:
            if result.column.id == "inserted_at":
                self.assertEqual(result.column.datatype, 'datetime')
                self.assertEqual(fake_time_now, result.value)
            elif result.column.id == "modified_on":
                self.assertEqual(result.column.datatype, 'datetime')
                self.assertEqual(modified_on, result.value)
            elif result.column.id == "opened_on":
                self.assertEqual(result.column.datatype, 'datetime')
                self.assertEqual(opened_on, result.value)
            elif result.column.id not in ["repeat_iteration", "inserted_at"]:
                self.assertEqual(sample_doc[_get_column_property(result.column)], result.value)
                if result.column.id in default_case_property_datatypes:
                    self.assertEqual(
                        result.column.datatype,
                        default_case_property_datatypes[result.column.id]
                    )

    def test_simple_form_management(self):
        app = Application.wrap(self.get_json('simple_app.json'))
        self.assertEqual('userreports_test', app.domain)
        data_sources = get_form_data_sources(app)
        self.assertEqual(1, len(data_sources))
        data_source = data_sources['http://openrosa.org/formdesigner/AF6F83BA-09A9-4773-9177-AB51EA6CF802']
        for indicator in data_source.configured_indicators:
            self.assertIsNotNone(indicator)
<commit_msg>use json_format_datetime<commit_after>import json
import os
from datetime import datetime
from mock import patch
from django.test import SimpleTestCase
from corehq.apps.app_manager.models import Application
from corehq.apps.userreports.app_manager import get_case_data_sources, get_form_data_sources
from corehq.apps.userreports.reports.builder import DEFAULT_CASE_PROPERTY_DATATYPES
from dimagi.utils.parsing import json_format_datetime

class AppManagerDataSourceConfigTest(SimpleTestCase):

    def setUp(self):
        self.is_usercase_in_use_patch = patch('corehq.apps.app_manager.models.is_usercase_in_use')
        is_usercase_in_use_mock = self.is_usercase_in_use_patch.start()
        is_usercase_in_use_mock.return_value = False

    def tearDown(self):
        self.is_usercase_in_use_patch.stop()

    def get_json(self, name):
        with open(os.path.join(os.path.dirname(__file__), 'data', 'app_manager', name)) as f:
            return json.loads(f.read())

    @patch('corehq.apps.userreports.specs.datetime')
    def test_simple_case_management(self, datetime_mock):
        fake_time_now = datetime(2015, 4, 24, 12, 30, 8, 24886)
        datetime_mock.utcnow.return_value = fake_time_now

        app = Application.wrap(self.get_json('simple_app.json'))
        self.assertEqual('userreports_test', app.domain)
        data_sources = get_case_data_sources(app)
        self.assertEqual(1, len(data_sources))
        data_source = data_sources['ticket']
        self.assertEqual(app.domain, data_source.domain)
        self.assertEqual('CommCareCase', data_source.referenced_doc_type)
        self.assertEqual('ticket', data_source.table_id)
        self.assertEqual('ticket', data_source.display_name)

        # test the filter
        self.assertTrue(data_source.filter(
            {'doc_type': 'CommCareCase', 'domain': app.domain, 'type': 'ticket'}))
        self.assertFalse(data_source.filter(
            {'doc_type': 'CommCareCase', 'domain': 'wrong domain', 'type': 'ticket'}))
        self.assertFalse(data_source.filter(
            {'doc_type': 'NotCommCareCase', 'domain': app.domain, 'type': 'ticket'}))
        self.assertFalse(data_source.filter(
            {'doc_type': 'CommCareCase', 'domain': app.domain, 'type': 'not-ticket'}))

        # check the indicators
        expected_columns = set(
            ["doc_id", "modified_on", "user_id", "opened_on",
             "owner_id", "name", "category", "priority", "starred", "estimate", "inserted_at"]
        )
        self.assertEqual(expected_columns, set(col_back.id for col_back in data_source.get_columns()))

        modified_on = datetime(2014, 11, 12, 15, 37, 49)
        opened_on = datetime(2014, 11, 11, 23, 34, 34, 25)
        sample_doc = dict(
            _id='some-doc-id',
            doc_type="CommCareCase",
            modified_on=json_format_datetime(modified_on),
            user_id="23407238074",
            opened_on=json_format_datetime(opened_on),
            owner_id="0923409230948",
            name="priority ticket",
            domain=app.domain,
            type='ticket',
            category='bug',
            priority='4',
            starred='yes',
            estimate='2',
        )

        def _get_column_property(column):
            return column.id if column.id != 'doc_id' else '_id'

        default_case_property_datatypes = DEFAULT_CASE_PROPERTY_DATATYPES
        [row] = data_source.get_all_values(sample_doc)
        for result in row:
            if result.column.id == "inserted_at":
                self.assertEqual(result.column.datatype, 'datetime')
                self.assertEqual(fake_time_now, result.value)
            elif result.column.id == "modified_on":
                self.assertEqual(result.column.datatype, 'datetime')
                self.assertEqual(modified_on, result.value)
            elif result.column.id == "opened_on":
                self.assertEqual(result.column.datatype, 'datetime')
                self.assertEqual(opened_on, result.value)
            elif result.column.id not in ["repeat_iteration", "inserted_at"]:
                self.assertEqual(sample_doc[_get_column_property(result.column)], result.value)
                if result.column.id in default_case_property_datatypes:
                    self.assertEqual(
                        result.column.datatype,
                        default_case_property_datatypes[result.column.id]
                    )

    def test_simple_form_management(self):
        app = Application.wrap(self.get_json('simple_app.json'))
        self.assertEqual('userreports_test', app.domain)
        data_sources = get_form_data_sources(app)
        self.assertEqual(1, len(data_sources))
        data_source = data_sources['http://openrosa.org/formdesigner/AF6F83BA-09A9-4773-9177-AB51EA6CF802']
        for indicator in data_source.configured_indicators:
            self.assertIsNotNone(indicator)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from __future__ import absolute_import

import os

try:
    import ConfigParser as configparser
except ImportError:
    # Python 3
    import configparser

__author__ = 'lundberg'

class IniConfigParser(object):

    section = 'main'

    def __init__(self, config_file_name, config_environment_variable=None):
        """
        :param config_file_name: config files name
        :param config_environment_variable: Optional env variable to read config from

        :type config_file_name: str | unicode
        :type config_environment_variable: str | unicode

        :return: IniConfigParser object
        :rtype: IniConfigParser
        """

        self.config_file_name = config_file_name
        self.config_environment_variable = config_environment_variable
        self.known_special_keys = {
            # <var name>.lower(): (<read function>, <default>),
            'celery_accept_content': (self.read_list, ['application/json']),
        }

    def get_config_file(self):
        """Get the configuration file looking for it in several places.

        The lookup order is:
        1. A file named according to the value of the config_environment_variable env variable
        2. A file named config_file_name In the current working directory
        3. A file named .config_file_name in the user's home directory
        4. A file named config_file_name in the system configuration directory (/etc)
        """
        file_name = os.environ.get(self.config_environment_variable, self.config_file_name)

        if os.path.exists(file_name):
            return file_name

        user_file = os.path.expanduser(
            os.path.join('~', '.', self.config_file_name))
        if os.path.exists(user_file):
            return user_file

        global_file = os.path.join('/etc', self.config_file_name)
        if os.path.exists(global_file):
            return global_file

    def read_setting_from_env(self, settings, key, default=None):
        env_variable = key.upper()
        if env_variable in os.environ:
            return os.environ[env_variable]
        else:
            return settings.get(key, default)

    def read_setting_from_env_bool(self, settings, key, default=None):
        value = self.read_setting_from_env(settings, key, '').lower()
        if value == 'false':
            return False
        if value == 'true':
            return True
        return default

    def read_setting_from_env_int(self, settings, key, default=None):
        value = self.read_setting_from_env(settings, key, '').lower()
        return int(value)

    def read_mapping(self, settings, prop, available_keys=None, default=None, required=True):
        raw = self.read_setting_from_env(settings, prop, '')

        if raw.strip() == '':
            return default

        rows = raw.strip('\n ').split('\n')

        mapping = {}

        for row in rows:
            splitted_row = row.split('=')
            key = splitted_row[0].strip()
            if len(splitted_row) > 1:
                value = splitted_row[1].strip()
            else:
                value = ''
            if available_keys is None or key in available_keys:
                mapping[key] = value

        if available_keys is not None:
            if len(mapping.keys()) != len(available_keys) and 'testing' not in settings:
                return None

        return mapping

    def read_list(self, settings, prop, default=list()):
        raw = self.read_setting_from_env(settings, prop, None)
        if raw is None or raw.strip() == '':
            return default

        return [e for e in raw.split('\n') if e is not None and e.strip() != '']

    def read_configuration(self):
        """
        Read the settings from environment or .ini file and return them as a dict

        The values are decoded as JSON, or used as-is if they can't be decoded as JSON.
        """
        settings = {}

        config_file = self.get_config_file()
        if config_file is not None:
            config = configparser.RawConfigParser()
            config.read(config_file)

            if config.has_section(self.section):
                for key, val in config.items(self.section):
                    if key in self.known_special_keys:
                        func, default = self.known_special_keys[key]
                    else:
                        func = self.read_setting_from_env
                        default = ''
                    settings[key.upper()] = func({key: val}, key, default=default)

        return settings<commit_msg>Missing newline at eof<commit_after># -*- coding: utf-8 -*-

from __future__ import absolute_import

import os

try:
    import ConfigParser as configparser
except ImportError:
    # Python 3
    import configparser

__author__ = 'lundberg'

class IniConfigParser(object):

    section = 'main'

    def __init__(self, config_file_name, config_environment_variable=None):
        """
        :param config_file_name: config files name
        :param config_environment_variable: Optional env variable to read config from

        :type config_file_name: str | unicode
        :type config_environment_variable: str | unicode

        :return: IniConfigParser object
        :rtype: IniConfigParser
        """

        self.config_file_name = config_file_name
        self.config_environment_variable = config_environment_variable
        self.known_special_keys = {
            # <var name>.lower(): (<read function>, <default>),
            'celery_accept_content': (self.read_list, ['application/json']),
        }

    def get_config_file(self):
        """Get the configuration file looking for it in several places.

        The lookup order is:
        1. A file named according to the value of the config_environment_variable env variable
        2. A file named config_file_name In the current working directory
        3. A file named .config_file_name in the user's home directory
        4. A file named config_file_name in the system configuration directory (/etc)
        """
        file_name = os.environ.get(self.config_environment_variable, self.config_file_name)

        if os.path.exists(file_name):
            return file_name

        user_file = os.path.expanduser(
            os.path.join('~', '.', self.config_file_name))
        if os.path.exists(user_file):
            return user_file

        global_file = os.path.join('/etc', self.config_file_name)
        if os.path.exists(global_file):
            return global_file

    def read_setting_from_env(self, settings, key, default=None):
        env_variable = key.upper()
        if env_variable in os.environ:
            return os.environ[env_variable]
        else:
            return settings.get(key, default)

    def read_setting_from_env_bool(self, settings, key, default=None):
        value = self.read_setting_from_env(settings, key, '').lower()
        if value == 'false':
            return False
        if value == 'true':
            return True
        return default

    def read_setting_from_env_int(self, settings, key, default=None):
        value = self.read_setting_from_env(settings, key, '').lower()
        return int(value)

    def read_mapping(self, settings, prop, available_keys=None, default=None, required=True):
        raw = self.read_setting_from_env(settings, prop, '')

        if raw.strip() == '':
            return default

        rows = raw.strip('\n ').split('\n')

        mapping = {}

        for row in rows:
            splitted_row = row.split('=')
            key = splitted_row[0].strip()
            if len(splitted_row) > 1:
                value = splitted_row[1].strip()
            else:
                value = ''
            if available_keys is None or key in available_keys:
                mapping[key] = value

        if available_keys is not None:
            if len(mapping.keys()) != len(available_keys) and 'testing' not in settings:
                return None

        return mapping

    def read_list(self, settings, prop, default=list()):
        raw = self.read_setting_from_env(settings, prop, None)
        if raw is None or raw.strip() == '':
            return default

        return [e for e in raw.split('\n') if e is not None and e.strip() != '']

    def read_configuration(self):
        """
        Read the settings from environment or .ini file and return them as a dict

        The values are decoded as JSON, or used as-is if they can't be decoded as JSON.
        """
        settings = {}

        config_file = self.get_config_file()
        if config_file is not None:
            config = configparser.RawConfigParser()
            config.read(config_file)

            if config.has_section(self.section):
                for key, val in config.items(self.section):
                    if key in self.known_special_keys:
                        func, default = self.known_special_keys[key]
                    else:
                        func = self.read_setting_from_env
                        default = ''
                    settings[key.upper()] = func({key: val}, key, default=default)

        return settings
<|endoftext|>


<commit_before>import asyncio
import os
import tempfile
import typing
from typing import Text, Optional, List, Union, Dict

from rasa import model, data
from rasa.cli.utils import (
    create_output_path,
    print_success,
    missing_config_keys,
    print_warning,
    get_validated_path,
)
from rasa.constants import (
    DEFAULT_MODELS_PATH,
    CONFIG_MANDATORY_KEYS,
    CONFIG_MANDATORY_KEYS_CORE,
    CONFIG_MANDATORY_KEYS_NLU,
    FALLBACK_CONFIG_PATH,
)

if typing.TYPE_CHECKING:
    from rasa.nlu.model import Interpreter

def train(
    domain: Text,
    config: Text,
    training_files: Union[Text, List[Text]],
    output: Text = DEFAULT_MODELS_PATH,
    force_training: bool = False,
    kwargs: Optional[Dict] = None,
) -> Optional[Text]:
    config = get_valid_config(config, CONFIG_MANDATORY_KEYS)

    loop = asyncio.get_event_loop()
    return loop.run_until_complete(
        train_async(domain, config, training_files, output, force_training, kwargs)
    )

async def train_async(
    domain: Text,
    config: Text,
    training_files: Union[Text, List[Text]],
    output: Text = DEFAULT_MODELS_PATH,
    force_training: bool = False,
    kwargs: Optional[Dict] = None,
) -> Optional[Text]:
    """Trains a Rasa model (Core and NLU).

    Args:
        domain: Path to the domain file.
        config: Path to the config for Core and NLU.
        training_files: Paths to the training data for Core and NLU.
        output: Output path.
        force_training: If `True` retrain model even if data has not changed.
        kwargs: Additional training parameters.

    Returns:
        Path of the trained model archive.
    """
    config = get_valid_config(config, CONFIG_MANDATORY_KEYS)

    train_path = tempfile.mkdtemp()
    old_model = model.get_latest_model(output)
    retrain_core = True
    retrain_nlu = True

    story_directory, nlu_data_directory = data.get_core_nlu_directories(training_files)
    new_fingerprint = model.model_fingerprint(
        config, domain, nlu_data_directory, story_directory
    )

    if not force_training and old_model:
        unpacked = model.unpack_model(old_model)
        old_core, old_nlu = model.get_model_subdirectories(unpacked)
        last_fingerprint = model.fingerprint_from_path(unpacked)

        if not model.core_fingerprint_changed(last_fingerprint, new_fingerprint):
            target_path = os.path.join(train_path, "core")
            retrain_core = not model.merge_model(old_core, target_path)

        if not model.nlu_fingerprint_changed(last_fingerprint, new_fingerprint):
            target_path = os.path.join(train_path, "nlu")
            retrain_nlu = not model.merge_model(old_nlu, target_path)

    if force_training or retrain_core:
        await train_core_async(
            domain, config, story_directory, output, train_path, kwargs
        )
    else:
        print (
            "Dialogue data / configuration did not change. "
            "No need to retrain dialogue model."
        )

    if force_training or retrain_nlu:
        train_nlu(config, nlu_data_directory, output, train_path)
    else:
        print ("NLU data / configuration did not change. No need to retrain NLU model.")

    if retrain_core or retrain_nlu:
        output = create_output_path(output)
        model.create_package_rasa(train_path, output, new_fingerprint)

        print ("Train path: '{}'.".format(train_path))

        print_success("Your bot is trained and ready to take for a spin!")

        return output
    else:
        print (
            "Nothing changed. You can use the old model stored at {}"
            "".format(os.path.abspath(old_model))
        )

        return old_model

def train_core(
    domain: Text,
    config: Text,
    stories: Text,
    output: Text,
    train_path: Optional[Text],
    kwargs: Optional[Dict],
) -> Optional[Text]:
    loop = asyncio.get_event_loop()
    return loop.run_until_complete(
        train_core_async(domain, config, stories, output, train_path, kwargs)
    )

async def train_core_async(
    domain: Text,
    config: Text,
    stories: Text,
    output: Text,
    train_path: Optional[Text] = None,
    kwargs: Optional[Dict] = None,
) -> Optional[Text]:
    """Trains a Core model.

    Args:
        domain: Path to the domain file.
        config: Path to the config file for Core.
        stories: Path to the Core training data.
        output: Output path.
        train_path: If `None` the model will be trained in a temporary
            directory, otherwise in the provided directory.
        kwargs: Additional training parameters.

    Returns:
        If `train_path` is given it returns the path to the model archive,
        otherwise the path to the directory with the trained model files.

    """
    import rasa.core.train

    config = get_valid_config(config, CONFIG_MANDATORY_KEYS_CORE)

    _train_path = train_path or tempfile.mkdtemp()

    # normal (not compare) training
    core_model = await rasa.core.train(
        domain_file=domain,
        stories_file=data.get_core_directory(stories),
        output_path=os.path.join(_train_path, "core"),
        policy_config=config,
        kwargs=kwargs,
    )

    if not train_path:
        # Only Core was trained.
        stories = data.get_core_directory(stories)
        output_path = create_output_path(output, prefix="core-")
        new_fingerprint = model.model_fingerprint(config, domain, stories=stories)
        model.create_package_rasa(_train_path, output_path, new_fingerprint)
        print_success(
            "Your Rasa Core model is trained and saved at '{}'.".format(output_path)
        )

    return core_model

def train_nlu(
    config: Text, nlu_data: Text, output: Text, train_path: Optional[Text]
) -> Optional["Interpreter"]:
    """Trains a NLU model.

    Args:
        config: Path to the config file for NLU.
        nlu_data: Path to the NLU training data.
        output: Output path.
        train_path: If `None` the model will be trained in a temporary
            directory, otherwise in the provided directory.

    Returns:
        If `train_path` is given it returns the path to the model archive,
        otherwise the path to the directory with the trained model files.

    """
    import rasa.nlu

    config = get_valid_config(config, CONFIG_MANDATORY_KEYS_NLU)

    _train_path = train_path or tempfile.mkdtemp()
    _, nlu_model, _ = rasa.nlu.train(
        config, data.get_nlu_directory(nlu_data), _train_path, fixed_model_name="nlu"
    )

    if not train_path:
        nlu_data = data.get_nlu_directory(nlu_data)
        output_path = create_output_path(output, prefix="nlu-")
        new_fingerprint = model.model_fingerprint(config, nlu_data=nlu_data)
        model.create_package_rasa(_train_path, output_path, new_fingerprint)
        print_success(
            "Your Rasa NLU model is trained and saved at '{}'.".format(output_path)
        )

    return nlu_model

def enrich_config(config_path, missing_keys, FALLBACK_CONFIG_PATH):
    import rasa.utils.io

    config_data = rasa.utils.io.read_yaml_file(config_path)
    fallback_config_data = rasa.utils.io.read_yaml_file(FALLBACK_CONFIG_PATH)

    for k in missing_keys:
        config_data[k] = fallback_config_data[k]

    rasa.utils.io.write_yaml_file(config_data, config_path)

def get_valid_config(config: Text, mandatory_keys: List[Text]) -> Text:
    config_path = get_validated_path(config, "config", FALLBACK_CONFIG_PATH)

    missing_keys = missing_config_keys(config_path, mandatory_keys)

    if missing_keys:
        print_warning(
            "Configuration file '{}' is missing mandatory parameters: "
            "{}. Filling missing parameters from fallback configuration file: '{}'."
            "".format(config, ", ".join(missing_keys), FALLBACK_CONFIG_PATH)
        )
        enrich_config(config_path, missing_keys, FALLBACK_CONFIG_PATH)

    return config_path
<commit_msg>Check if nlu argument is a directory<commit_after>import asyncio
import os
import tempfile
import typing
from typing import Text, Optional, List, Union, Dict

from rasa import model, data
from rasa.cli.utils import (
    create_output_path,
    print_success,
    missing_config_keys,
    print_warning,
    get_validated_path,
)
from rasa.constants import (
    DEFAULT_MODELS_PATH,
    CONFIG_MANDATORY_KEYS,
    CONFIG_MANDATORY_KEYS_CORE,
    CONFIG_MANDATORY_KEYS_NLU,
    FALLBACK_CONFIG_PATH,
)

if typing.TYPE_CHECKING:
    from rasa.nlu.model import Interpreter

def train(
    domain: Text,
    config: Text,
    training_files: Union[Text, List[Text]],
    output: Text = DEFAULT_MODELS_PATH,
    force_training: bool = False,
    kwargs: Optional[Dict] = None,
) -> Optional[Text]:
    config = get_valid_config(config, CONFIG_MANDATORY_KEYS)

    loop = asyncio.get_event_loop()
    return loop.run_until_complete(
        train_async(domain, config, training_files, output, force_training, kwargs)
    )

async def train_async(
    domain: Text,
    config: Text,
    training_files: Union[Text, List[Text]],
    output: Text = DEFAULT_MODELS_PATH,
    force_training: bool = False,
    kwargs: Optional[Dict] = None,
) -> Optional[Text]:
    """Trains a Rasa model (Core and NLU).

    Args:
        domain: Path to the domain file.
        config: Path to the config for Core and NLU.
        training_files: Paths to the training data for Core and NLU.
        output: Output path.
        force_training: If `True` retrain model even if data has not changed.
        kwargs: Additional training parameters.

    Returns:
        Path of the trained model archive.
    """
    config = get_valid_config(config, CONFIG_MANDATORY_KEYS)

    train_path = tempfile.mkdtemp()
    old_model = model.get_latest_model(output)
    retrain_core = True
    retrain_nlu = True

    story_directory, nlu_data_directory = data.get_core_nlu_directories(training_files)
    new_fingerprint = model.model_fingerprint(
        config, domain, nlu_data_directory, story_directory
    )

    if not force_training and old_model:
        unpacked = model.unpack_model(old_model)
        old_core, old_nlu = model.get_model_subdirectories(unpacked)
        last_fingerprint = model.fingerprint_from_path(unpacked)

        if not model.core_fingerprint_changed(last_fingerprint, new_fingerprint):
            target_path = os.path.join(train_path, "core")
            retrain_core = not model.merge_model(old_core, target_path)

        if not model.nlu_fingerprint_changed(last_fingerprint, new_fingerprint):
            target_path = os.path.join(train_path, "nlu")
            retrain_nlu = not model.merge_model(old_nlu, target_path)

    if force_training or retrain_core:
        await train_core_async(
            domain, config, story_directory, output, train_path, kwargs
        )
    else:
        print (
            "Dialogue data / configuration did not change. "
            "No need to retrain dialogue model."
        )

    if force_training or retrain_nlu:
        train_nlu(config, nlu_data_directory, output, train_path)
    else:
        print ("NLU data / configuration did not change. No need to retrain NLU model.")

    if retrain_core or retrain_nlu:
        output = create_output_path(output)
        model.create_package_rasa(train_path, output, new_fingerprint)

        print ("Train path: '{}'.".format(train_path))

        print_success("Your bot is trained and ready to take for a spin!")

        return output
    else:
        print (
            "Nothing changed. You can use the old model stored at {}"
            "".format(os.path.abspath(old_model))
        )

        return old_model

def train_core(
    domain: Text,
    config: Text,
    stories: Text,
    output: Text,
    train_path: Optional[Text],
    kwargs: Optional[Dict],
) -> Optional[Text]:
    loop = asyncio.get_event_loop()
    return loop.run_until_complete(
        train_core_async(domain, config, stories, output, train_path, kwargs)
    )

async def train_core_async(
    domain: Text,
    config: Text,
    stories: Text,
    output: Text,
    train_path: Optional[Text] = None,
    kwargs: Optional[Dict] = None,
) -> Optional[Text]:
    """Trains a Core model.

    Args:
        domain: Path to the domain file.
        config: Path to the config file for Core.
        stories: Path to the Core training data.
        output: Output path.
        train_path: If `None` the model will be trained in a temporary
            directory, otherwise in the provided directory.
        kwargs: Additional training parameters.

    Returns:
        If `train_path` is given it returns the path to the model archive,
        otherwise the path to the directory with the trained model files.

    """
    import rasa.core.train

    config = get_valid_config(config, CONFIG_MANDATORY_KEYS_CORE)

    _train_path = train_path or tempfile.mkdtemp()

    # normal (not compare) training
    core_model = await rasa.core.train(
        domain_file=domain,
        stories_file=data.get_core_directory(stories),
        output_path=os.path.join(_train_path, "core"),
        policy_config=config,
        kwargs=kwargs,
    )

    if not train_path:
        # Only Core was trained.
        stories = data.get_core_directory(stories)
        output_path = create_output_path(output, prefix="core-")
        new_fingerprint = model.model_fingerprint(config, domain, stories=stories)
        model.create_package_rasa(_train_path, output_path, new_fingerprint)
        print_success(
            "Your Rasa Core model is trained and saved at '{}'.".format(output_path)
        )

    return core_model

def train_nlu(
    config: Text, nlu_data: Text, output: Text, train_path: Optional[Text]
) -> Optional["Interpreter"]:
    """Trains a NLU model.

    Args:
        config: Path to the config file for NLU.
        nlu_data: Path to the NLU training data.
        output: Output path.
        train_path: If `None` the model will be trained in a temporary
            directory, otherwise in the provided directory.

    Returns:
        If `train_path` is given it returns the path to the model archive,
        otherwise the path to the directory with the trained model files.

    """
    import rasa.nlu

    config = get_valid_config(config, CONFIG_MANDATORY_KEYS_NLU)

    if os.path.isdir(nlu_data):
        nlu_data = data.get_nlu_directory(nlu_data)

    _train_path = train_path or tempfile.mkdtemp()
    _, nlu_model, _ = rasa.nlu.train(
        config, nlu_data, _train_path, fixed_model_name="nlu"
    )

    if not train_path:
        output_path = create_output_path(output, prefix="nlu-")
        new_fingerprint = model.model_fingerprint(config, nlu_data=nlu_data)
        model.create_package_rasa(_train_path, output_path, new_fingerprint)
        print_success(
            "Your Rasa NLU model is trained and saved at '{}'.".format(output_path)
        )

    return nlu_model

def enrich_config(config_path, missing_keys, FALLBACK_CONFIG_PATH):
    import rasa.utils.io

    config_data = rasa.utils.io.read_yaml_file(config_path)
    fallback_config_data = rasa.utils.io.read_yaml_file(FALLBACK_CONFIG_PATH)

    for k in missing_keys:
        config_data[k] = fallback_config_data[k]

    rasa.utils.io.write_yaml_file(config_data, config_path)

def get_valid_config(config: Text, mandatory_keys: List[Text]) -> Text:
    config_path = get_validated_path(config, "config", FALLBACK_CONFIG_PATH)

    missing_keys = missing_config_keys(config_path, mandatory_keys)

    if missing_keys:
        print_warning(
            "Configuration file '{}' is missing mandatory parameters: "
            "{}. Filling missing parameters from fallback configuration file: '{}'."
            "".format(config, ", ".join(missing_keys), FALLBACK_CONFIG_PATH)
        )
        enrich_config(config_path, missing_keys, FALLBACK_CONFIG_PATH)

    return config_path
<|endoftext|>