<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import io
import json
import logging
import os
import re
import warnings
from typing import Optional, List, Text, Any, Dict, AnyStr

from rasa_core import utils
from rasa_core.events import (
    ActionExecuted, UserUttered, Event, SlotSet)
from rasa_core.exceptions import StoryParseError
from rasa_core.interpreter import RegexInterpreter
from rasa_core.training.structures import (
    Checkpoint, STORY_START, StoryStep,
    GENERATED_CHECKPOINT_PREFIX, GENERATED_HASH_LENGTH, FORM_PREFIX)
from rasa_nlu import utils as nlu_utils
from rasa_nlu.training_data import Message
from rasa_nlu.training_data.formats import MarkdownReader

logger = logging.getLogger(__name__)

class EndToEndReader(MarkdownReader):
    def _parse_item(self, line):
        # type: (Text) -> Optional[Message]
        """Parses an md list item line based on the current section type.

        Matches expressions of the form `<intent>:<example>. For the
        syntax of <example> see the Rasa NLU docs on training data:
        https://rasa.com/docs/nlu/dataformat/#markdown-format"""

        item_regex = re.compile(r'\s*(.+?):\s*(.*)')
        match = re.match(item_regex, line)
        if match:
            intent = match.group(1)
            self.current_title = intent
            message = match.group(2)
            example = self._parse_training_example(message)
            example.data["true_intent"] = intent
            return example

        raise ValueError("Encountered invalid end-to-end format for message "
                         "`{}`. Please visit the documentation page on "
                         "end-to-end evaluation at https://rasa.com/docs/core/"
                         "evaluation#end-to-end-evaluation-of-rasa-nlu-and-"
                         "core".format(line))

class StoryStepBuilder(object):
    def __init__(self, name):
        self.name = name
        self.story_steps = []
        self.current_steps = []
        self.start_checkpoints = []

    def add_checkpoint(self, name, conditions):
        # type: (Text, Optional[Dict[Text, Any]]) -> None

        # Depending on the state of the story part this
        # is either a start or an end check point
        if not self.current_steps:
            self.start_checkpoints.append(Checkpoint(name, conditions))
        else:
            if conditions:
                logger.warning("End or intermediate checkpoints "
                               "do not support conditions! "
                               "(checkpoint: {})".format(name))
            additional_steps = []
            for t in self.current_steps:
                if t.end_checkpoints:
                    tcp = t.create_copy(use_new_id=True)
                    tcp.end_checkpoints = [Checkpoint(name)]
                    additional_steps.append(tcp)
                else:
                    t.end_checkpoints = [Checkpoint(name)]
            self.current_steps.extend(additional_steps)

    def _prev_end_checkpoints(self):
        if not self.current_steps:
            return self.start_checkpoints
        else:
            # makes sure we got each end name only once
            end_names = {e.name
                         for s in self.current_steps
                         for e in s.end_checkpoints}
            return [Checkpoint(name) for name in end_names]

    def add_user_messages(self, messages):
        self.ensure_current_steps()

        if len(messages) == 1:
            # If there is only one possible intent, we'll keep things simple
            for t in self.current_steps:
                t.add_user_message(messages[0])
        else:
            # If there are multiple different intents the
            # user can use the express the same thing
            # we need to copy the blocks and create one
            # copy for each possible message
            prefix = GENERATED_CHECKPOINT_PREFIX + "OR_"
            generated_checkpoint = utils.generate_id(prefix,
                                                     GENERATED_HASH_LENGTH)
            updated_steps = []
            for t in self.current_steps:
                for m in messages:
                    copied = t.create_copy(use_new_id=True)
                    copied.add_user_message(m)
                    copied.end_checkpoints = [Checkpoint(generated_checkpoint)]
                    updated_steps.append(copied)
            self.current_steps = updated_steps

    def add_event(self, event):
        self.ensure_current_steps()
        for t in self.current_steps:
            t.add_event(event)

    def ensure_current_steps(self):
        completed = [step
                     for step in self.current_steps
                     if step.end_checkpoints]
        unfinished = [step
                      for step in self.current_steps
                      if not step.end_checkpoints]
        self.story_steps.extend(completed)
        if unfinished:
            self.current_steps = unfinished
        else:
            self.current_steps = self._next_story_steps()

    def flush(self):
        if self.current_steps:
            self.story_steps.extend(self.current_steps)
            self.current_steps = []

    def _next_story_steps(self):
        start_checkpoints = self._prev_end_checkpoints()
        if not start_checkpoints:
            start_checkpoints = [Checkpoint(STORY_START)]
        current_turns = [StoryStep(block_name=self.name,
                                   start_checkpoints=start_checkpoints)]
        return current_turns

class StoryFileReader(object):
    """Helper class to read a story file."""

    def __init__(self, domain, interpreter, template_vars=None, use_e2e=False):
        self.story_steps = []
        self.current_step_builder = None  # type: Optional[StoryStepBuilder]
        self.domain = domain
        self.interpreter = interpreter
        self.template_variables = template_vars if template_vars else {}
        self.use_e2e = use_e2e

    @staticmethod
    def read_from_folder(resource_name, domain, interpreter=RegexInterpreter(),
                         template_variables=None, use_e2e=False,
                         exclusion_percentage=None):
        """Given a path reads all contained story files."""

        story_steps = []
        for f in nlu_utils.list_files(resource_name):
            steps = StoryFileReader.read_from_file(f, domain, interpreter,
                                                   template_variables, use_e2e)
            story_steps.extend(steps)

            # if exclusion percentage is not 100
            if exclusion_percentage and exclusion_percentage is not 100:
                import random
                idx = int(round(exclusion_percentage/100.0 * len(story_steps)))
                random.shuffle(story_steps)
                story_steps = story_steps[:-idx]

        return story_steps

    @staticmethod
    def read_from_file(filename, domain, interpreter=RegexInterpreter(),
                       template_variables=None, use_e2e=False):
        """Given a md file reads the contained stories."""

        try:
            with io.open(filename, "r", encoding="utf-8") as f:
                lines = f.readlines()
            reader = StoryFileReader(domain, interpreter,
                                     template_variables, use_e2e)
            return reader.process_lines(lines)
        except ValueError as err:
            file_info = ("Invalid story file format. Failed to parse "
                         "'{}'".format(os.path.abspath(filename)))
            logger.exception(file_info)
            if not err.args:
                err.args = ('',)
            err.args = err.args + (file_info,)
            raise

    @staticmethod
    def _parameters_from_json_string(s, line):
        # type: (Text, Text) -> Dict[Text, Any]
        """Parse the passed string as json and create a parameter dict."""

        if s is None or not s.strip():
            # if there is no strings there are not going to be any parameters
            return {}

        try:
            parsed_slots = json.loads(s)
            if isinstance(parsed_slots, dict):
                return parsed_slots
            else:
                raise Exception("Parsed value isn't a json object "
                                "(instead parser found '{}')"
                                ".".format(type(parsed_slots)))
        except Exception as e:
            raise ValueError("Invalid to parse arguments in line "
                             "'{}'. Failed to decode parameters"
                             "as a json object. Make sure the event"
                             "name is followed by a proper json "
                             "object. Error: {}".format(line, e))

    @staticmethod
    def _parse_event_line(line):
        """Tries to parse a single line as an event with arguments."""

        # the regex matches "slot{"a": 1}"
        m = re.search('^([^{]+)([{].+)?', line)
        if m is not None:
            event_name = m.group(1).strip()
            slots_str = m.group(2)
            parameters = StoryFileReader._parameters_from_json_string(
                    slots_str, line)
            return event_name, parameters
        else:
            warnings.warn("Failed to parse action line '{}'. "
                          "Ignoring this line.".format(line))
            return "", {}

    def process_lines(self, lines):
        # type: (List[AnyStr]) -> List[StoryStep]

        for idx, line in enumerate(lines):
            line_num = idx + 1
            try:
                line = self._replace_template_variables(
                        self._clean_up_line(line))
                if line.strip() == "":
                    continue
                elif line.startswith("#"):  # reached a new story block
                    name = line[1:].strip("# ")
                    self.new_story_part(name)
                elif line.startswith(">"):  # reached a checkpoint
                    name, conditions = self._parse_event_line(line[1:].strip())
                    self.add_checkpoint(name, conditions)
                elif re.match(r'^[*\-]\s+{}'.format(FORM_PREFIX), line):
                    logger.debug("Skipping line {}, "
                                 "because it was generated by "
                                 "form action".format(line))
                elif line.startswith(
                        "-"):  # reached a slot, event, or executed action
                    event_name, parameters = self._parse_event_line(line[1:])
                    self.add_event(event_name, parameters)
                elif line.startswith("*"):  # reached a user message
                    user_messages = [el.strip() for el in
                                     line[1:].split(" OR ")]
                    if self.use_e2e:
                        self.add_e2e_messages(user_messages, line_num)
                    else:
                        self.add_user_messages(user_messages, line_num)
                else:  # reached an unknown type of line
                    logger.warning("Skipping line {}. "
                                   "No valid command found. "
                                   "Line Content: '{}'"
                                   "".format(line_num, line))
            except Exception as e:
                msg = "Error in line {}: {}".format(line_num, e)
                logger.error(msg, exc_info=1)
                raise ValueError(msg)
        self._add_current_stories_to_result()
        return self.story_steps

    def _replace_template_variables(self, line):
        def process_match(matchobject):
            varname = matchobject.group(1)
            if varname in self.template_variables:
                return self.template_variables[varname]
            else:
                raise ValueError("Unknown variable `{var}` "
                                 "in template line '{line}'"
                                 "".format(var=varname, line=line))

        template_rx = re.compile(r"`([^`]+)`")
        return template_rx.sub(process_match, line)

    @staticmethod
    def _clean_up_line(line):
        # type: (Text) -> Text
        """Removes comments and trailing spaces"""

        return re.sub(r'<!--.*?-->', '', line).strip()

    def _add_current_stories_to_result(self):
        if self.current_step_builder:
            self.current_step_builder.flush()
            self.story_steps.extend(self.current_step_builder.story_steps)

    def new_story_part(self, name):
        self._add_current_stories_to_result()
        self.current_step_builder = StoryStepBuilder(name)

    def add_checkpoint(self, name, conditions):
        # type: (Text, Optional[Dict[Text, Any]]) -> None

        # Ensure story part already has a name
        if not self.current_step_builder:
            raise StoryParseError("Checkpoint '{}' is at an invalid location. "
                                  "Expected a story start.".format(name))

        self.current_step_builder.add_checkpoint(name, conditions)

    def _parse_message(self, message, line_num):
        parse_data = self.interpreter.parse(message)
        utterance = UserUttered(message,
                                parse_data.get("intent"),
                                parse_data.get("entities"),
                                parse_data)

        if message.startswith("_"):
            c = utterance.as_story_string()
            logger.warning("Stating user intents with a leading '_' is "
                           "deprecated. The new format is "
                           "'* {}'. Please update "
                           "your example '{}' to the new format."
                           "".format(c, message))
        intent_name = utterance.intent.get("name")
        if intent_name not in self.domain.intents:
            logger.warning("Found unknown intent '{}' on line {}. "
                           "Please, make sure that all intents are "
                           "listed in your domain yaml."
                           "".format(intent_name, line_num))
        return utterance

    def add_user_messages(self, messages, line_num):
        if not self.current_step_builder:
            raise StoryParseError("User message '{}' at invalid location. "
                                  "Expected story start.".format(messages))
        parsed_messages = [self._parse_message(m, line_num) for m in messages]
        self.current_step_builder.add_user_messages(parsed_messages)

    def add_e2e_messages(self, e2e_messages, line_num):
        if not self.current_step_builder:
            raise StoryParseError("End-to-end message '{}' at invalid "
                                  "location. Expected story start."
                                  "".format(e2e_messages))
        e2e_reader = EndToEndReader()
        parsed_messages = []
        for m in e2e_messages:
            message = e2e_reader._parse_item(m)
            parsed = self._parse_message(message.text, line_num)

            parsed.parse_data["true_intent"] = message.data["true_intent"]
            parsed.parse_data["true_entities"] = \
                message.data.get("entities") or []
            parsed_messages.append(parsed)
        self.current_step_builder.add_user_messages(parsed_messages)

    def add_event(self, event_name, parameters):

        # add 'name' only if event is not a SlotSet,
        # because there might be a slot with slot_key='name'
        if "name" not in parameters and event_name != SlotSet.type_name:
            parameters["name"] = event_name

        parsed_events = Event.from_story_string(event_name, parameters,
                                                default=ActionExecuted)
        if parsed_events is None:
            raise StoryParseError("Unknown event '{}'. It is Neither an event "
                                  "nor an action).".format(event_name))
        if self.current_step_builder is None:
            raise StoryParseError("Failed to handle event '{}'. There is no "
                                  "started story block available. "
                                  "".format(event_name))

        for p in parsed_events:
            self.current_step_builder.add_event(p)
<commit_msg>fix indent<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import io
import json
import logging
import os
import re
import warnings
from typing import Optional, List, Text, Any, Dict, AnyStr

from rasa_core import utils
from rasa_core.events import (
    ActionExecuted, UserUttered, Event, SlotSet)
from rasa_core.exceptions import StoryParseError
from rasa_core.interpreter import RegexInterpreter
from rasa_core.training.structures import (
    Checkpoint, STORY_START, StoryStep,
    GENERATED_CHECKPOINT_PREFIX, GENERATED_HASH_LENGTH, FORM_PREFIX)
from rasa_nlu import utils as nlu_utils
from rasa_nlu.training_data import Message
from rasa_nlu.training_data.formats import MarkdownReader

logger = logging.getLogger(__name__)

class EndToEndReader(MarkdownReader):
    def _parse_item(self, line):
        # type: (Text) -> Optional[Message]
        """Parses an md list item line based on the current section type.

        Matches expressions of the form `<intent>:<example>. For the
        syntax of <example> see the Rasa NLU docs on training data:
        https://rasa.com/docs/nlu/dataformat/#markdown-format"""

        item_regex = re.compile(r'\s*(.+?):\s*(.*)')
        match = re.match(item_regex, line)
        if match:
            intent = match.group(1)
            self.current_title = intent
            message = match.group(2)
            example = self._parse_training_example(message)
            example.data["true_intent"] = intent
            return example

        raise ValueError("Encountered invalid end-to-end format for message "
                         "`{}`. Please visit the documentation page on "
                         "end-to-end evaluation at https://rasa.com/docs/core/"
                         "evaluation#end-to-end-evaluation-of-rasa-nlu-and-"
                         "core".format(line))

class StoryStepBuilder(object):
    def __init__(self, name):
        self.name = name
        self.story_steps = []
        self.current_steps = []
        self.start_checkpoints = []

    def add_checkpoint(self, name, conditions):
        # type: (Text, Optional[Dict[Text, Any]]) -> None

        # Depending on the state of the story part this
        # is either a start or an end check point
        if not self.current_steps:
            self.start_checkpoints.append(Checkpoint(name, conditions))
        else:
            if conditions:
                logger.warning("End or intermediate checkpoints "
                               "do not support conditions! "
                               "(checkpoint: {})".format(name))
            additional_steps = []
            for t in self.current_steps:
                if t.end_checkpoints:
                    tcp = t.create_copy(use_new_id=True)
                    tcp.end_checkpoints = [Checkpoint(name)]
                    additional_steps.append(tcp)
                else:
                    t.end_checkpoints = [Checkpoint(name)]
            self.current_steps.extend(additional_steps)

    def _prev_end_checkpoints(self):
        if not self.current_steps:
            return self.start_checkpoints
        else:
            # makes sure we got each end name only once
            end_names = {e.name
                         for s in self.current_steps
                         for e in s.end_checkpoints}
            return [Checkpoint(name) for name in end_names]

    def add_user_messages(self, messages):
        self.ensure_current_steps()

        if len(messages) == 1:
            # If there is only one possible intent, we'll keep things simple
            for t in self.current_steps:
                t.add_user_message(messages[0])
        else:
            # If there are multiple different intents the
            # user can use the express the same thing
            # we need to copy the blocks and create one
            # copy for each possible message
            prefix = GENERATED_CHECKPOINT_PREFIX + "OR_"
            generated_checkpoint = utils.generate_id(prefix,
                                                     GENERATED_HASH_LENGTH)
            updated_steps = []
            for t in self.current_steps:
                for m in messages:
                    copied = t.create_copy(use_new_id=True)
                    copied.add_user_message(m)
                    copied.end_checkpoints = [Checkpoint(generated_checkpoint)]
                    updated_steps.append(copied)
            self.current_steps = updated_steps

    def add_event(self, event):
        self.ensure_current_steps()
        for t in self.current_steps:
            t.add_event(event)

    def ensure_current_steps(self):
        completed = [step
                     for step in self.current_steps
                     if step.end_checkpoints]
        unfinished = [step
                      for step in self.current_steps
                      if not step.end_checkpoints]
        self.story_steps.extend(completed)
        if unfinished:
            self.current_steps = unfinished
        else:
            self.current_steps = self._next_story_steps()

    def flush(self):
        if self.current_steps:
            self.story_steps.extend(self.current_steps)
            self.current_steps = []

    def _next_story_steps(self):
        start_checkpoints = self._prev_end_checkpoints()
        if not start_checkpoints:
            start_checkpoints = [Checkpoint(STORY_START)]
        current_turns = [StoryStep(block_name=self.name,
                                   start_checkpoints=start_checkpoints)]
        return current_turns

class StoryFileReader(object):
    """Helper class to read a story file."""

    def __init__(self, domain, interpreter, template_vars=None, use_e2e=False):
        self.story_steps = []
        self.current_step_builder = None  # type: Optional[StoryStepBuilder]
        self.domain = domain
        self.interpreter = interpreter
        self.template_variables = template_vars if template_vars else {}
        self.use_e2e = use_e2e

    @staticmethod
    def read_from_folder(resource_name, domain, interpreter=RegexInterpreter(),
                         template_variables=None, use_e2e=False,
                         exclusion_percentage=None):
        """Given a path reads all contained story files."""

        story_steps = []
        for f in nlu_utils.list_files(resource_name):
            steps = StoryFileReader.read_from_file(f, domain, interpreter,
                                                   template_variables, use_e2e)
            story_steps.extend(steps)

        # if exclusion percentage is not 100
        if exclusion_percentage and exclusion_percentage is not 100:
            import random
            idx = int(round(exclusion_percentage/100.0 * len(story_steps)))
            random.shuffle(story_steps)
            story_steps = story_steps[:-idx]

        return story_steps

    @staticmethod
    def read_from_file(filename, domain, interpreter=RegexInterpreter(),
                       template_variables=None, use_e2e=False):
        """Given a md file reads the contained stories."""

        try:
            with io.open(filename, "r", encoding="utf-8") as f:
                lines = f.readlines()
            reader = StoryFileReader(domain, interpreter,
                                     template_variables, use_e2e)
            return reader.process_lines(lines)
        except ValueError as err:
            file_info = ("Invalid story file format. Failed to parse "
                         "'{}'".format(os.path.abspath(filename)))
            logger.exception(file_info)
            if not err.args:
                err.args = ('',)
            err.args = err.args + (file_info,)
            raise

    @staticmethod
    def _parameters_from_json_string(s, line):
        # type: (Text, Text) -> Dict[Text, Any]
        """Parse the passed string as json and create a parameter dict."""

        if s is None or not s.strip():
            # if there is no strings there are not going to be any parameters
            return {}

        try:
            parsed_slots = json.loads(s)
            if isinstance(parsed_slots, dict):
                return parsed_slots
            else:
                raise Exception("Parsed value isn't a json object "
                                "(instead parser found '{}')"
                                ".".format(type(parsed_slots)))
        except Exception as e:
            raise ValueError("Invalid to parse arguments in line "
                             "'{}'. Failed to decode parameters"
                             "as a json object. Make sure the event"
                             "name is followed by a proper json "
                             "object. Error: {}".format(line, e))

    @staticmethod
    def _parse_event_line(line):
        """Tries to parse a single line as an event with arguments."""

        # the regex matches "slot{"a": 1}"
        m = re.search('^([^{]+)([{].+)?', line)
        if m is not None:
            event_name = m.group(1).strip()
            slots_str = m.group(2)
            parameters = StoryFileReader._parameters_from_json_string(
                    slots_str, line)
            return event_name, parameters
        else:
            warnings.warn("Failed to parse action line '{}'. "
                          "Ignoring this line.".format(line))
            return "", {}

    def process_lines(self, lines):
        # type: (List[AnyStr]) -> List[StoryStep]

        for idx, line in enumerate(lines):
            line_num = idx + 1
            try:
                line = self._replace_template_variables(
                        self._clean_up_line(line))
                if line.strip() == "":
                    continue
                elif line.startswith("#"):  # reached a new story block
                    name = line[1:].strip("# ")
                    self.new_story_part(name)
                elif line.startswith(">"):  # reached a checkpoint
                    name, conditions = self._parse_event_line(line[1:].strip())
                    self.add_checkpoint(name, conditions)
                elif re.match(r'^[*\-]\s+{}'.format(FORM_PREFIX), line):
                    logger.debug("Skipping line {}, "
                                 "because it was generated by "
                                 "form action".format(line))
                elif line.startswith(
                        "-"):  # reached a slot, event, or executed action
                    event_name, parameters = self._parse_event_line(line[1:])
                    self.add_event(event_name, parameters)
                elif line.startswith("*"):  # reached a user message
                    user_messages = [el.strip() for el in
                                     line[1:].split(" OR ")]
                    if self.use_e2e:
                        self.add_e2e_messages(user_messages, line_num)
                    else:
                        self.add_user_messages(user_messages, line_num)
                else:  # reached an unknown type of line
                    logger.warning("Skipping line {}. "
                                   "No valid command found. "
                                   "Line Content: '{}'"
                                   "".format(line_num, line))
            except Exception as e:
                msg = "Error in line {}: {}".format(line_num, e)
                logger.error(msg, exc_info=1)
                raise ValueError(msg)
        self._add_current_stories_to_result()
        return self.story_steps

    def _replace_template_variables(self, line):
        def process_match(matchobject):
            varname = matchobject.group(1)
            if varname in self.template_variables:
                return self.template_variables[varname]
            else:
                raise ValueError("Unknown variable `{var}` "
                                 "in template line '{line}'"
                                 "".format(var=varname, line=line))

        template_rx = re.compile(r"`([^`]+)`")
        return template_rx.sub(process_match, line)

    @staticmethod
    def _clean_up_line(line):
        # type: (Text) -> Text
        """Removes comments and trailing spaces"""

        return re.sub(r'<!--.*?-->', '', line).strip()

    def _add_current_stories_to_result(self):
        if self.current_step_builder:
            self.current_step_builder.flush()
            self.story_steps.extend(self.current_step_builder.story_steps)

    def new_story_part(self, name):
        self._add_current_stories_to_result()
        self.current_step_builder = StoryStepBuilder(name)

    def add_checkpoint(self, name, conditions):
        # type: (Text, Optional[Dict[Text, Any]]) -> None

        # Ensure story part already has a name
        if not self.current_step_builder:
            raise StoryParseError("Checkpoint '{}' is at an invalid location. "
                                  "Expected a story start.".format(name))

        self.current_step_builder.add_checkpoint(name, conditions)

    def _parse_message(self, message, line_num):
        parse_data = self.interpreter.parse(message)
        utterance = UserUttered(message,
                                parse_data.get("intent"),
                                parse_data.get("entities"),
                                parse_data)

        if message.startswith("_"):
            c = utterance.as_story_string()
            logger.warning("Stating user intents with a leading '_' is "
                           "deprecated. The new format is "
                           "'* {}'. Please update "
                           "your example '{}' to the new format."
                           "".format(c, message))
        intent_name = utterance.intent.get("name")
        if intent_name not in self.domain.intents:
            logger.warning("Found unknown intent '{}' on line {}. "
                           "Please, make sure that all intents are "
                           "listed in your domain yaml."
                           "".format(intent_name, line_num))
        return utterance

    def add_user_messages(self, messages, line_num):
        if not self.current_step_builder:
            raise StoryParseError("User message '{}' at invalid location. "
                                  "Expected story start.".format(messages))
        parsed_messages = [self._parse_message(m, line_num) for m in messages]
        self.current_step_builder.add_user_messages(parsed_messages)

    def add_e2e_messages(self, e2e_messages, line_num):
        if not self.current_step_builder:
            raise StoryParseError("End-to-end message '{}' at invalid "
                                  "location. Expected story start."
                                  "".format(e2e_messages))
        e2e_reader = EndToEndReader()
        parsed_messages = []
        for m in e2e_messages:
            message = e2e_reader._parse_item(m)
            parsed = self._parse_message(message.text, line_num)

            parsed.parse_data["true_intent"] = message.data["true_intent"]
            parsed.parse_data["true_entities"] = \
                message.data.get("entities") or []
            parsed_messages.append(parsed)
        self.current_step_builder.add_user_messages(parsed_messages)

    def add_event(self, event_name, parameters):

        # add 'name' only if event is not a SlotSet,
        # because there might be a slot with slot_key='name'
        if "name" not in parameters and event_name != SlotSet.type_name:
            parameters["name"] = event_name

        parsed_events = Event.from_story_string(event_name, parameters,
                                                default=ActionExecuted)
        if parsed_events is None:
            raise StoryParseError("Unknown event '{}'. It is Neither an event "
                                  "nor an action).".format(event_name))
        if self.current_step_builder is None:
            raise StoryParseError("Failed to handle event '{}'. There is no "
                                  "started story block available. "
                                  "".format(event_name))

        for p in parsed_events:
            self.current_step_builder.add_event(p)
<|endoftext|>


<commit_before>"""
Tests for the mock of the database summary endpoint.
"""

from time import sleep

import pytest
import requests
import timeout_decorator
from requests import codes
from requests_mock import GET

from common.constants import ResultCodes
from tests.mock_vws.utils import assert_vws_response
from tests.utils import VuforiaServerCredentials
from vws._request_utils import target_api_request

def database_summary(
    vuforia_server_credentials: VuforiaServerCredentials,
) -> requests.Response:
    """
    Return the response of a request to the database summary endpoint.

    Args:
        vuforia_server_credentials: The credentials to use to connect to
            Vuforia.
    """
    return target_api_request(
        access_key=vuforia_server_credentials.access_key,
        secret_key=vuforia_server_credentials.secret_key,
        method=GET,
        content=b'',
        request_path='/summary',
    )

@timeout_decorator.timeout(seconds=200)
def wait_for_image_numbers(
    vuforia_server_credentials: VuforiaServerCredentials,
    active_images: int,
    inactive_images: int,
    failed_images: int,
    processing_images: int,
) -> None:
    """
    Wait up to 200 seconds (arbitrary) for the number of images in various
    categories to match the expected number.

    This is necessary because the database summary endpoint lags behind the
    real data.

    This is susceptible to false positives because if, for example, we expect
    no images, and the endpoint adds images with a delay, we will not know.

    Args:
        vuforia_server_credentials: The credentials to use to connect to
            Vuforia.
        active_images: The expected number of active images.
        inactive_images: The expected number of inactive images.
        failed_images: The expected number of failed images.
        processing_images: The expected number of processing images.

    Raises:
        TimeoutError: The numbers of images in various categories do not match
            within the time limit.
    """
    while True:
        response = database_summary(
            vuforia_server_credentials=vuforia_server_credentials
        )

        if all(
            [
                active_images == response.json()['active_images'],
                inactive_images == response.json()['inactive_images'],
                failed_images == response.json()['failed_images'],
                processing_images == response.json()['processing_images'],
            ]
        ):
            return

        # We wait 0.2 seconds rather than less than that to decrease the number
        # of calls made to the API, to decrease the likelihood of hitting the
        # request quota.
        sleep(0.2)

@pytest.mark.usefixtures('verify_mock_vuforia')
class TestDatabaseSummary:
    """
    Tests for the mock of the database summary endpoint at `GET /summary`.
    """

    def test_success(
        self,
        vuforia_server_credentials: VuforiaServerCredentials,
    ) -> None:
        """
        It is possible to get a success response.
        """
        response = database_summary(
            vuforia_server_credentials=vuforia_server_credentials
        )

        assert_vws_response(
            response=response,
            status_code=codes.OK,
            result_code=ResultCodes.SUCCESS,
        )

        assert response.json().keys() == {
            'active_images',
            'current_month_recos',
            'failed_images',
            'inactive_images',
            'name',
            'previous_month_recos',
            'processing_images',
            'reco_threshold',
            'request_quota',
            'request_usage',
            'result_code',
            'target_quota',
            'total_recos',
            'transaction_id',
        }

        assert response.json()['name'] == (
            vuforia_server_credentials.database_name
        )

        wait_for_image_numbers(
            vuforia_server_credentials=vuforia_server_credentials,
            active_images=0,
            inactive_images=0,
            failed_images=0,
            processing_images=0,

@pytest.mark.usefixtures('verify_mock_vuforia_inactive')
class TestInactiveProject:
    """
    Tests for inactive projects.
    """

    def test_inactive_project(
        self,
        inactive_server_credentials: VuforiaServerCredentials,
    ) -> None:
        """
        The project's active state does not affect the database summary.
        """
        response = database_summary(
            vuforia_server_credentials=inactive_server_credentials,
        )

        assert_vws_response(
            response=response,
            status_code=codes.OK,
            result_code=ResultCodes.SUCCESS,
        )
<commit_msg>Fix merge error<commit_after>"""
Tests for the mock of the database summary endpoint.
"""

from time import sleep

import pytest
import requests
import timeout_decorator
from requests import codes
from requests_mock import GET

from common.constants import ResultCodes
from tests.mock_vws.utils import assert_vws_response
from tests.utils import VuforiaServerCredentials
from vws._request_utils import target_api_request

def database_summary(
    vuforia_server_credentials: VuforiaServerCredentials,
) -> requests.Response:
    """
    Return the response of a request to the database summary endpoint.

    Args:
        vuforia_server_credentials: The credentials to use to connect to
            Vuforia.
    """
    return target_api_request(
        access_key=vuforia_server_credentials.access_key,
        secret_key=vuforia_server_credentials.secret_key,
        method=GET,
        content=b'',
        request_path='/summary',
    )

@timeout_decorator.timeout(seconds=200)
def wait_for_image_numbers(
    vuforia_server_credentials: VuforiaServerCredentials,
    active_images: int,
    inactive_images: int,
    failed_images: int,
    processing_images: int,
) -> None:
    """
    Wait up to 200 seconds (arbitrary) for the number of images in various
    categories to match the expected number.

    This is necessary because the database summary endpoint lags behind the
    real data.

    This is susceptible to false positives because if, for example, we expect
    no images, and the endpoint adds images with a delay, we will not know.

    Args:
        vuforia_server_credentials: The credentials to use to connect to
            Vuforia.
        active_images: The expected number of active images.
        inactive_images: The expected number of inactive images.
        failed_images: The expected number of failed images.
        processing_images: The expected number of processing images.

    Raises:
        TimeoutError: The numbers of images in various categories do not match
            within the time limit.
    """
    while True:
        response = database_summary(
            vuforia_server_credentials=vuforia_server_credentials
        )

        if all(
            [
                active_images == response.json()['active_images'],
                inactive_images == response.json()['inactive_images'],
                failed_images == response.json()['failed_images'],
                processing_images == response.json()['processing_images'],
            ]
        ):
            return

        # We wait 0.2 seconds rather than less than that to decrease the number
        # of calls made to the API, to decrease the likelihood of hitting the
        # request quota.
        sleep(0.2)

@pytest.mark.usefixtures('verify_mock_vuforia')
class TestDatabaseSummary:
    """
    Tests for the mock of the database summary endpoint at `GET /summary`.
    """

    def test_success(
        self,
        vuforia_server_credentials: VuforiaServerCredentials,
    ) -> None:
        """
        It is possible to get a success response.
        """
        response = database_summary(
            vuforia_server_credentials=vuforia_server_credentials
        )

        assert_vws_response(
            response=response,
            status_code=codes.OK,
            result_code=ResultCodes.SUCCESS,
        )

        assert response.json().keys() == {
            'active_images',
            'current_month_recos',
            'failed_images',
            'inactive_images',
            'name',
            'previous_month_recos',
            'processing_images',
            'reco_threshold',
            'request_quota',
            'request_usage',
            'result_code',
            'target_quota',
            'total_recos',
            'transaction_id',
        }

        assert response.json()['name'] == (
            vuforia_server_credentials.database_name
        )

        wait_for_image_numbers(
            vuforia_server_credentials=vuforia_server_credentials,
            active_images=0,
            inactive_images=0,
            failed_images=0,
            processing_images=0,
        )

@pytest.mark.usefixtures('verify_mock_vuforia_inactive')
class TestInactiveProject:
    """
    Tests for inactive projects.
    """

    def test_inactive_project(
        self,
        inactive_server_credentials: VuforiaServerCredentials,
    ) -> None:
        """
        The project's active state does not affect the database summary.
        """
        response = database_summary(
            vuforia_server_credentials=inactive_server_credentials,
        )

        assert_vws_response(
            response=response,
            status_code=codes.OK,
            result_code=ResultCodes.SUCCESS,
        )
<|endoftext|>


<commit_before>"""
Tests for passing invalid endpoints which require a target ID to be given.
"""

import uuid
from urllib.parse import urljoin

import pytest
import requests
from requests import codes
from requests_mock import GET

from common.constants import ResultCodes
from tests.conftest import VuforiaServerCredentials
from tests.mock_vws.utils import assert_vws_failure
from vws._request_utils import authorization_header, rfc_1123_date

class Endpoint:
    """
    XXX
    """

    def __init__(self, path: str, method: int) -> None:
        """
        """
        self.path = path
        self.method = method

@pytest.fixture()
def target_list():
    return Endpoint(path='/targets', method=GET)

@pytest.fixture()
def get_duplicates():
    return Endpoint(path='/duplicates', method=GET)

@pytest.fixture(params=['target_list', 'get_duplicates'])
def endpoint(request):
    return request.getfixturevalue(request.param)

@pytest.mark.usefixtures('verify_mock_vuforia')
class TestInvalidGivenID:
    """
    Tests for giving an invalid ID to endpoints which require a target ID to
    be given.
    """

    def test_not_real_id(self,
                         vuforia_server_credentials: VuforiaServerCredentials,
                         endpoint: Endpoint,
                         ) -> None:
        """
        A `NOT_FOUND` error is returned when an endpoint is given a target ID
        of a target which does not exist.
        """
        request_path = endpoint.path + '/' + uuid.uuid4().hex
        date = rfc_1123_date()

        authorization_string = authorization_header(
            access_key=vuforia_server_credentials.access_key,
            secret_key=vuforia_server_credentials.secret_key,
            method=GET,
            content=b'',
            content_type='',
            date=date,
            request_path=request_path,
        )

        headers = {
            "Authorization": authorization_string,
            "Date": date,
        }

        url = urljoin('https://vws.vuforia.com/', request_path)
        response = requests.request(
            method=GET,
            url=url,
            headers=headers,
            data=b'',
        )

        assert_vws_failure(
            response=response,
            status_code=codes.NOT_FOUND,
            result_code=ResultCodes.UNKNOWN_TARGET,
        )
<commit_msg>Type new fixtures<commit_after>"""
Tests for passing invalid endpoints which require a target ID to be given.
"""

import uuid
from urllib.parse import urljoin

import pytest
import requests
from _pytest.fixtures import SubRequest
from requests import codes
from requests_mock import GET

from common.constants import ResultCodes
from tests.conftest import VuforiaServerCredentials
from tests.mock_vws.utils import assert_vws_failure
from vws._request_utils import authorization_header, rfc_1123_date

class Endpoint:
    """
    Details of endpoints to be called in tests.
    """

    def __init__(self, path: str, method: int) -> None:
        """
        Args:
            path: X
            method: X
        """
        self.path = path
        self.method = method

@pytest.fixture()
def target_list() -> Endpoint:
    """
    XXX
    """
    return Endpoint(path='/targets', method=GET)

@pytest.fixture()
def get_duplicates() -> Endpoint:
    """
    XXX
    """
    return Endpoint(path='/duplicates', method=GET)

@pytest.fixture(params=['target_list', 'get_duplicates'])
def endpoint(request: SubRequest) -> Endpoint:
    """
    XXX
    """
    return request.getfixturevalue(request.param)

@pytest.mark.usefixtures('verify_mock_vuforia')
class TestInvalidGivenID:
    """
    Tests for giving an invalid ID to endpoints which require a target ID to
    be given.
    """

    def test_not_real_id(self,
                         vuforia_server_credentials: VuforiaServerCredentials,
                         endpoint: Endpoint,
                         ) -> None:
        """
        A `NOT_FOUND` error is returned when an endpoint is given a target ID
        of a target which does not exist.
        """
        request_path = endpoint.path + '/' + uuid.uuid4().hex
        date = rfc_1123_date()

        authorization_string = authorization_header(
            access_key=vuforia_server_credentials.access_key,
            secret_key=vuforia_server_credentials.secret_key,
            method=GET,
            content=b'',
            content_type='',
            date=date,
            request_path=request_path,
        )

        headers = {
            "Authorization": authorization_string,
            "Date": date,
        }

        url = urljoin('https://vws.vuforia.com/', request_path)
        response = requests.request(
            method=GET,
            url=url,
            headers=headers,
            data=b'',
        )

        assert_vws_failure(
            response=response,
            status_code=codes.NOT_FOUND,
            result_code=ResultCodes.UNKNOWN_TARGET,
        )
<|endoftext|>


<commit_before>from __future__ import generators

ANY = None

from rdflib.backends import Backend

class Memory(Backend):
    """\
An in memory implementation of a triple store.

This triple store uses nested dictionaries to store triples. Each
triple is stored in two such indices as follows spo[s][p][o] = 1 and
pos[p][o][s] = 1.
    """
    def __init__(self):
        super(Memory, self).__init__()
        self.context_aware = False
        # indexed by [subject][predicate][object]
        self.__spo = {}

        # indexed by [predicate][object][subject]
        self.__pos = {}

        # indexed by [predicate][object][subject]
        self.__osp = {}

        self.__namespace = {}
        self.__prefix = {}

    def add(self, (subject, predicate, object), context=None, quoted=False):
        """\
        Add a triple to the store of triples.
        """
        # add dictionary entries for spo[s][p][p] = 1 and pos[p][o][s]
        # = 1, creating the nested dictionaries where they do not yet
        # exits.
        spo = self.__spo
        try:
            po = spo[subject]
        except:
            po = spo[subject] = {}
        try:
            o = po[predicate]
        except:
            o = po[predicate] = {}
        o[object] = 1

        pos = self.__pos
        try:
            os = pos[predicate]
        except:
            os = pos[predicate] = {}
        try:
            s = os[object]
        except:
            s = os[object] = {}
        s[subject] = 1

        osp = self.__osp
        try:
            sp = osp[object]
        except:
            sp = osp[object] = {}
        try:
            p = sp[subject]
        except:
            p = sp[subject] = {}
        p[predicate] = 1

    def remove(self, (subject, predicate, object), context=None):
        for subject, predicate, object in self.triples((subject, predicate, object)):
            del self.__spo[subject][predicate][object]
            del self.__pos[predicate][object][subject]
            del self.__osp[object][subject][predicate]

    def triples(self, (subject, predicate, object), context=None):
        """A generator over all the triples matching """
        if subject!=ANY: # subject is given
            spo = self.__spo
            if subject in spo:
                subjectDictionary = spo[subject]
                if predicate!=ANY: # subject+predicate is given
                    if predicate in subjectDictionary:
                        if object!=ANY: # subject+predicate+object is given
                            if object in subjectDictionary[predicate]:
                                yield (subject, predicate, object)
                            else: # given object not found
                                pass
                        else: # subject+predicate is given, object unbound
                            for o in subjectDictionary[predicate].keys():
                                yield (subject, predicate, o)
                    else: # given predicate not found
                        pass
                else: # subject given, predicate unbound
                    for p in subjectDictionary.keys():
                        if object!=ANY: # object is given
                            if object in subjectDictionary[p]:
                                yield (subject, p, object)
                            else: # given object not found
                                pass
                        else: # object unbound
                            for o in subjectDictionary[p].keys():
                                yield (subject, p, o)
            else: # given subject not found
                pass
        elif predicate!=ANY: # predicate is given, subject unbound
            pos = self.__pos
            if predicate in pos:
                predicateDictionary = pos[predicate]
                if object!=ANY: # predicate+object is given, subject unbound
                    if object in predicateDictionary:
                        for s in predicateDictionary[object].keys():
                            yield (s, predicate, object)
                    else: # given object not found
                        pass
                else: # predicate is given, object+subject unbound
                    for o in predicateDictionary.keys():
                        for s in predicateDictionary[o].keys():
                            yield (s, predicate, o)
        elif object!=ANY: # object is given, subject+predicate unbound
            osp = self.__osp
            if object in osp:
                objectDictionary = osp[object]
                for s in objectDictionary.keys():
                    for p in objectDictionary[s].keys():
                        yield (s, p, object)
        else: # subject+predicate+object unbound
            spo = self.__spo
            for s in spo.keys():
                subjectDictionary = spo[s]
                for p in subjectDictionary.keys():
                    for o in subjectDictionary[p].keys():
                        yield (s, p, o)

    def __len__(self, context=None):
        #@@ optimize
        i = 0
        for triple in self.triples((None, None, None)):
            i += 1
        return i

    def bind(self, prefix, namespace):
        self.__prefix[namespace] = prefix
        self.__namespace[prefix] = namespace

    def namespace(self, prefix):
        return self.__namespace.get(prefix, None)

    def prefix(self, namespace):
        return self.__prefix.get(namespace, None)

    def namespaces(self):
        for prefix, namespace in self.__namespace.iteritems():
            yield prefix, namespace
<commit_msg>updating to store interface<commit_after>from __future__ import generators

ANY = None

from rdflib.stores import Store

class Memory(Store):
    """\
An in memory implementation of a triple store.

This triple store uses nested dictionaries to store triples. Each
triple is stored in two such indices as follows spo[s][p][o] = 1 and
pos[p][o][s] = 1.
    """
    def __init__(self, identifier, configuration=None):
        super(Memory, self).__init__(identifier, configuration)
        # indexed by [subject][predicate][object]
        self.__spo = {}

        # indexed by [predicate][object][subject]
        self.__pos = {}

        # indexed by [predicate][object][subject]
        self.__osp = {}

        self.__namespace = {}
        self.__prefix = {}

    def add(self, (subject, predicate, object), context=None, quoted=False):
        """\
        Add a triple to the store of triples.
        """
        # add dictionary entries for spo[s][p][p] = 1 and pos[p][o][s]
        # = 1, creating the nested dictionaries where they do not yet
        # exits.
        spo = self.__spo
        try:
            po = spo[subject]
        except:
            po = spo[subject] = {}
        try:
            o = po[predicate]
        except:
            o = po[predicate] = {}
        o[object] = 1

        pos = self.__pos
        try:
            os = pos[predicate]
        except:
            os = pos[predicate] = {}
        try:
            s = os[object]
        except:
            s = os[object] = {}
        s[subject] = 1

        osp = self.__osp
        try:
            sp = osp[object]
        except:
            sp = osp[object] = {}
        try:
            p = sp[subject]
        except:
            p = sp[subject] = {}
        p[predicate] = 1

    def remove(self, (subject, predicate, object), context=None):
        for subject, predicate, object in self.triples((subject, predicate, object)):
            del self.__spo[subject][predicate][object]
            del self.__pos[predicate][object][subject]
            del self.__osp[object][subject][predicate]

    def triples(self, (subject, predicate, object), context=None):
        """A generator over all the triples matching """
        if subject!=ANY: # subject is given
            spo = self.__spo
            if subject in spo:
                subjectDictionary = spo[subject]
                if predicate!=ANY: # subject+predicate is given
                    if predicate in subjectDictionary:
                        if object!=ANY: # subject+predicate+object is given
                            if object in subjectDictionary[predicate]:
                                yield (subject, predicate, object), self.__contexts()
                            else: # given object not found
                                pass
                        else: # subject+predicate is given, object unbound
                            for o in subjectDictionary[predicate].keys():
                                yield (subject, predicate, o), self.__contexts()
                    else: # given predicate not found
                        pass
                else: # subject given, predicate unbound
                    for p in subjectDictionary.keys():
                        if object!=ANY: # object is given
                            if object in subjectDictionary[p]:
                                yield (subject, p, object), self.__contexts()
                            else: # given object not found
                                pass
                        else: # object unbound
                            for o in subjectDictionary[p].keys():
                                yield (subject, p, o), self.__contexts()
            else: # given subject not found
                pass
        elif predicate!=ANY: # predicate is given, subject unbound
            pos = self.__pos
            if predicate in pos:
                predicateDictionary = pos[predicate]
                if object!=ANY: # predicate+object is given, subject unbound
                    if object in predicateDictionary:
                        for s in predicateDictionary[object].keys():
                            yield (s, predicate, object), self.__contexts()
                    else: # given object not found
                        pass
                else: # predicate is given, object+subject unbound
                    for o in predicateDictionary.keys():
                        for s in predicateDictionary[o].keys():
                            yield (s, predicate, o), self.__contexts()
        elif object!=ANY: # object is given, subject+predicate unbound
            osp = self.__osp
            if object in osp:
                objectDictionary = osp[object]
                for s in objectDictionary.keys():
                    for p in objectDictionary[s].keys():
                        yield (s, p, object), self.__contexts()
        else: # subject+predicate+object unbound
            spo = self.__spo
            for s in spo.keys():
                subjectDictionary = spo[s]
                for p in subjectDictionary.keys():
                    for o in subjectDictionary[p].keys():
                        yield (s, p, o), self.__contexts()

    def __len__(self, context=None):
        #@@ optimize
        i = 0
        for triple in self.triples((None, None, None)):
            i += 1
        return i

    def bind(self, prefix, namespace):
        self.__prefix[namespace] = prefix
        self.__namespace[prefix] = namespace

    def namespace(self, prefix):
        return self.__namespace.get(prefix, None)

    def prefix(self, namespace):
        return self.__prefix.get(namespace, None)

    def namespaces(self):
        for prefix, namespace in self.__namespace.iteritems():
            yield prefix, namespace

    def __contexts(self):
        yield self.identifier

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function

from tests.test_api.utils import TestBaseApi

from polyaxon_client.api.base import BaseApiHandler
from polyaxon_client.exceptions import ERRORS_MAPPING

class DummyConfig(object):
    @staticmethod
    def from_dict(value):
        return 'from_dict', value

class TestBaseApiHandler(TestBaseApi):
    # pylint:disable=protected-access
    def setUp(self):
        super(TestBaseApiHandler, self).setUp()
        self.api_handler = BaseApiHandler(transport=self.transport, config=self.api_config)

    def test_get_page(self):
        assert self.api_handler.get_page() == {}
        assert self.api_handler.get_page(page=1) == {}
        assert self.api_handler.get_page(page=2) == {'offset': self.api_config.PAGE_SIZE}
        assert self.api_handler.get_page(page=3) == {'offset': self.api_config.PAGE_SIZE * 2}

    def test_build_url(self):
        assert self.api_handler.build_url('a') == 'a/'
        assert self.api_handler.build_url('a', 'b') == 'a/b/'

    def test_get_url(self):
        with self.assertRaises(ERRORS_MAPPING['base']):
            self.api_handler._get_url('a')

        assert self.api_handler._get_url('base', 'endpoint') == 'base/endpoint/'
        assert self.api_handler._get_http_url('endpoint') == '{}/endpoint/'.format(
            self.api_config.base_url)
        assert self.api_handler._get_ws_url('endpoint') == '{}/endpoint/'.format(
            self.api_config.base_ws_url)

    def test_prepare_results(self):
        # Schema response
        results = self.api_handler.prepare_results(response_json={'foo': 'bar'},
                                                   config=DummyConfig)
        assert results == ('from_dict', {'foo': 'bar'})

        # Raw response
        self.api_config.schema_response = False
        results = self.api_handler.prepare_results(response_json={'foo': 'bar'},
                                                   config=DummyConfig)
        assert results == {'foo': 'bar'}

    def test_prepare_list_results(self):
        # Schema response
        results = self.api_handler.prepare_list_results(response_json={'results': ['bar']},
                                                        current_page=1,
                                                        config=DummyConfig)
        assert results == {'count': 0,
                           'next': None,
                           'previous': None,
                           'results': [('from_dict', 'bar')]}

        # Raw response
        self.api_config.schema_response = False
        results = self.api_handler.prepare_list_results(response_json={'results': ['bar']},
                                                        current_page=1,
                                                        config=DummyConfig)
        assert results == {'count': 0, 'next': None, 'previous': None, 'results': ['bar']}
<commit_msg>Update tests<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import, division, print_function

from tests.test_api.utils import TestBaseApi

from polyaxon_client.api.base import BaseApiHandler
from polyaxon_client.exceptions import ERRORS_MAPPING

class DummyConfig(object):
    @staticmethod
    def from_dict(value, unknown=None):  # noqa
        return 'from_dict', value

class TestBaseApiHandler(TestBaseApi):
    # pylint:disable=protected-access
    def setUp(self):
        super(TestBaseApiHandler, self).setUp()
        self.api_handler = BaseApiHandler(transport=self.transport, config=self.api_config)

    def test_get_page(self):
        assert self.api_handler.get_page() == {}
        assert self.api_handler.get_page(page=1) == {}
        assert self.api_handler.get_page(page=2) == {'offset': self.api_config.PAGE_SIZE}
        assert self.api_handler.get_page(page=3) == {'offset': self.api_config.PAGE_SIZE * 2}

    def test_build_url(self):
        assert self.api_handler.build_url('a') == 'a/'
        assert self.api_handler.build_url('a', 'b') == 'a/b/'

    def test_get_url(self):
        with self.assertRaises(ERRORS_MAPPING['base']):
            self.api_handler._get_url('a')

        assert self.api_handler._get_url('base', 'endpoint') == 'base/endpoint/'
        assert self.api_handler._get_http_url('endpoint') == '{}/endpoint/'.format(
            self.api_config.base_url)
        assert self.api_handler._get_ws_url('endpoint') == '{}/endpoint/'.format(
            self.api_config.base_ws_url)

    def test_prepare_results(self):
        # Schema response
        results = self.api_handler.prepare_results(response_json={'foo': 'bar'},
                                                   config=DummyConfig)
        assert results == ('from_dict', {'foo': 'bar'})

        # Raw response
        self.api_config.schema_response = False
        results = self.api_handler.prepare_results(response_json={'foo': 'bar'},
                                                   config=DummyConfig)
        assert results == {'foo': 'bar'}

    def test_prepare_list_results(self):
        # Schema response
        results = self.api_handler.prepare_list_results(response_json={'results': ['bar']},
                                                        current_page=1,
                                                        config=DummyConfig)
        assert results == {'count': 0,
                           'next': None,
                           'previous': None,
                           'results': [('from_dict', 'bar')]}

        # Raw response
        self.api_config.schema_response = False
        results = self.api_handler.prepare_list_results(response_json={'results': ['bar']},
                                                        current_page=1,
                                                        config=DummyConfig)
        assert results == {'count': 0, 'next': None, 'previous': None, 'results': ['bar']}
<|endoftext|>


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

from __future__ import (absolute_import, division,
                        print_function)
from ..read_config import read_config
import numpy as np
import matplotlib.pyplot as plt
from ..imager import ScanSet, main_imager
from ..simulate import simulate_map
from ..global_fit import display_intermediate
from ..calibration import CalibratorTable
from ..io import mkdir_p
import os
import glob
import astropy.units as u

try:
    from tqdm import tqdm
except ImportError:
    def tqdm(x):
        return x

np.random.seed(1241347)

def _2d_gauss(x, y, sigma=3 / 60.):
    """A Gaussian beam"""
    return np.exp(-(x ** 2 + y ** 2) / (2 * sigma**2))

def gauss_src_func(x, y):
    return 50 * _2d_gauss(x, y, sigma=3 / 60)

def sim_config_file(filename):
    """Create a sample config file, to be modified by hand."""
    string = """
[local]
workdir : .
datadir : .

[analysis]
projection : ARC
interpolation : spline
prefix : test_
list_of_directories :
    gauss_ra
    gauss_dec
calibrator_directories :
    calibration

noise_threshold : 5

pixel_size : 0.8
        """
    with open(filename, 'w') as fobj:
        print(string, file=fobj)
    return string

def sim_map(obsdir_ra, obsdir_dec):
    simulate_map(count_map=gauss_src_func,
                 length_ra=50.,
                 length_dec=50.,
                 outdir=(obsdir_ra, obsdir_dec), mean_ra=180,
                 mean_dec=70, speed=2.,
                 spacing=0.5, srcname='Dummy')

class TestScanSet(object):
    @classmethod
    def setup_class(klass):
        import os
        global DEBUG_MODE
        DEBUG_MODE = True

        klass.curdir = os.path.dirname(__file__)
        klass.datadir = os.path.join(klass.curdir, 'data')
        klass.obsdir_ra = os.path.join(klass.datadir, 'sim', 'gauss_ra')
        klass.obsdir_dec = os.path.join(klass.datadir, 'sim', 'gauss_dec')
        klass.config_file = \
            os.path.abspath(os.path.join(klass.datadir, 'sim',
                                         'test_config_sim.ini'))
        klass.caldir = os.path.join(klass.datadir, 'sim', 'calibration')
        # First off, simulate a beamed observation  -------

        if not os.path.exists(klass.config_file):
            print('Setting up simulated data.')
            sim_config_file(klass.config_file)

        if (not os.path.exists(klass.obsdir_ra)) or \
                (not os.path.exists(klass.obsdir_dec)):
            mkdir_p(klass.obsdir_ra)
            mkdir_p(klass.obsdir_dec)
            print('Fake map: Point-like (but Gaussian beam shape), 0.5 Jy.')
            sim_map(klass.obsdir_ra, klass.obsdir_dec)

        caltable = CalibratorTable()
        caltable.from_scans(glob.glob(os.path.join(klass.caldir,
                                                   '*.fits')), debug=True)

        caltable.update()
        klass.calfile = os.path.join(klass.datadir, 'calibrators.hdf5')
        caltable.write(klass.calfile, overwrite=True)

        klass.config = read_config(klass.config_file)
        if os.path.exists('test.hdf5'):
            klass.scanset = ScanSet('test.hdf5')
        else:
            klass.scanset = ScanSet(klass.config_file, norefilt=False,
                                    debug=True)
            klass.scanset.write('test.hdf5', overwrite=True)

        plt.ioff()

    def test_0_prepare(self):
        pass

    def test_use_command_line(self):
        main_imager(('test.hdf5 -u Jy/beam '.format(self.config_file) +
                     '--calibrate {}'.format(self.calfile) +
                     ' -o bubu.hdf5 --debug').split(' '))

    def test_1_meta_saved_and_loaded_correctly(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)
        for k in scanset.meta.keys():
            assert np.all(scanset.meta[k] == self.scanset.meta[k])
        assert sorted(scanset.meta.keys()) == sorted(self.scanset.meta.keys())
        assert scanset.scan_list == self.scanset.scan_list

    def test_2_rough_image(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        images = scanset.calculate_images()

        img = images['Ch0']

        fig = plt.figure('img')
        plt.imshow(img, origin='lower')
        plt.colorbar()
        plt.savefig('img.png')
        plt.close(fig)

    def test_3_rough_image_altaz(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        images = scanset.calculate_images(altaz=True)

        img = images['Ch0']

        fig = plt.figure('img_altaz')
        plt.imshow(img, origin='lower')
        plt.colorbar()
        plt.savefig('img_altaz.png')
        scanset.save_ds9_images(save_sdev=True, altaz=True)
        plt.close(fig)

    def test_4_image_stdev(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        images = scanset.calculate_images()

        img = images['Ch0-Sdev']

        fig = plt.figure('log(img-Sdev)')
        plt.imshow(np.log10(img), origin='lower')
        plt.colorbar()
        plt.ioff()
        plt.savefig('img_sdev.png')
        plt.close(fig)

    def test_5_image_scrunch(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        images = scanset.calculate_images(scrunch=True)

        img = images['Ch0']

        fig = plt.figure('img - scrunched')
        plt.imshow(img, origin='lower')
        plt.colorbar()
        img = images['Ch0-Sdev']
        plt.savefig('img_scrunch.png')
        plt.close(fig)

        fig = plt.figure('img - scrunched - sdev')
        plt.imshow(img, origin='lower')
        plt.colorbar()
        plt.ioff()
        plt.savefig('img_scrunch_sdev.png')
        plt.close(fig)

    def test_6a_calibrate_image_pixel(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        images = scanset.calculate_images(calibration=self.calfile,
                                          map_unit="Jy/pixel")

        img = images['Ch0']
        center = img.shape[0] // 2, img.shape[1] // 2
        shortest_side = np.min(img.shape)
        X, Y = np.meshgrid(np.arange(img.shape[1]), np.arange(img.shape[0]))
        good = (X-center[1])**2 + (Y-center[0])**2 <= (shortest_side//4)**2
        assert np.allclose(np.sum(images['Ch0'][good]), 0.5, 0.05)

    def test_6b_calibrate_image_beam(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        images = scanset.calculate_images(calibration=self.calfile,
                                          map_unit="Jy/beam")

        assert np.allclose(np.max(images['Ch0']), 0.5, atol=0.05)

    def test_6c_calibrate_image_sr(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        images = scanset.calculate_images(calibration=self.calfile,
                                          map_unit="Jy/sr")

        images_pix = scanset.calculate_images(calibration=self.calfile,
                                              map_unit="Jy/pixel")

        pixel_area = scanset.meta['pixel_size'] ** 2
        assert np.allclose(images['Ch0'],
                           images_pix['Ch0'] / pixel_area.to(u.sr).value,
                           rtol=0.05)

    def test_7_ds9_image(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.save_ds9_images(save_sdev=True)

    def test_8_global_fit_image(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)
        images = scanset.calculate_images()
        nx, ny = images['Ch0'].shape
        excluded = [[nx//2, ny//2, nx//4]]
        scanset.fit_full_images(excluded=excluded, chans='Ch0')
        os.path.exists("out_iter_Ch0_002.txt")
        scanset.fit_full_images(excluded=excluded, chans='Ch1')
        os.path.exists("out_iter_Ch1_000.txt")
        display_intermediate(scanset, chan="Ch0", excluded=excluded,
                             parfile="out_iter_Ch0_002.txt")
        os.path.exists("out_iter_Ch1_002.png")

    def test_9a_find_scan_through_pixel(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        _, _, _, _, _, _, _, coord = \
            scanset.find_scans_through_pixel(62, 62, test=True)

        dec_scan = os.path.join(self.obsdir_dec, 'Dec99.fits')
        assert dec_scan in coord
        assert coord[dec_scan] == 'dec'
        ra_scan = os.path.join(self.obsdir_ra, 'Ra100.fits')
        assert ra_scan in coord
        assert coord[ra_scan] == 'ra'

    def test_9b_find_scan_through_pixel(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        _, _, _, _, _, _, _, coord = \
            scanset.find_scans_through_pixel(62, 0, test=True)

        dec_scan = os.path.join(self.obsdir_dec, 'Dec99.fits')
        assert dec_scan in coord
        assert coord[dec_scan] == 'dec'
        ra_scan = os.path.join(self.obsdir_ra, 'Ra0.fits')
        assert ra_scan in coord
        assert coord[ra_scan] == 'ra'

    def test_9c_find_scan_through_invalid_pixel(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        _, _, _, _, _, _, _, coord = \
            scanset.find_scans_through_pixel(62, -2, test=True)
        assert coord == {}
        _, _, _, _, _, _, _, coord = \
            scanset.find_scans_through_pixel(62, 64, test=True)
        assert coord == {}

    @classmethod
    def teardown_class(klass):
        """Clean up the mess."""
        os.unlink('img.png')
        os.unlink('img_altaz.png')
        os.unlink('img_scrunch.png')
        os.unlink('delta_altaz.png')
        os.unlink('altaz_with_src.png')
        os.unlink('img_sdev.png')
        os.unlink('test.hdf5')
        os.unlink('test_scan_list.txt')
        os.unlink('img_scrunch_sdev.png')
        os.unlink('bubu.hdf5')
        for d in klass.config['list_of_directories']:
            hfiles = \
                glob.glob(os.path.join(klass.config['datadir'], d, '*.hdf5'))
            for h in hfiles:
                os.unlink(h)
        out_iter_files = glob.glob('out_iter_*')
        for o in out_iter_files:
            os.unlink(o)
        out_fits_files = glob.glob(os.path.join(klass.config['datadir'],
                                                'test_config*.fits'))
        out_hdf5_files = glob.glob(os.path.join(klass.config['datadir'], 'sim',
                                                '*/', '*.hdf5'))

        for o in out_fits_files + out_hdf5_files:
            os.unlink(o)
<commit_msg>Cleanup [skip ci]<commit_after># -*- coding: utf-8 -*-

from __future__ import (absolute_import, division,
                        print_function)
from ..read_config import read_config
import numpy as np
import matplotlib.pyplot as plt
from ..imager import ScanSet, main_imager
from ..simulate import simulate_map
from ..global_fit import display_intermediate
from ..calibration import CalibratorTable
from ..io import mkdir_p
import os
import glob
import astropy.units as u

try:
    from tqdm import tqdm
except ImportError:
    def tqdm(x):
        return x

np.random.seed(1241347)

def _2d_gauss(x, y, sigma=3 / 60.):
    """A Gaussian beam"""
    return np.exp(-(x ** 2 + y ** 2) / (2 * sigma**2))

def gauss_src_func(x, y):
    return 50 * _2d_gauss(x, y, sigma=3 / 60)

def sim_config_file(filename):
    """Create a sample config file, to be modified by hand."""
    string = """
[local]
workdir : .
datadir : .

[analysis]
projection : ARC
interpolation : spline
prefix : test_
list_of_directories :
    gauss_ra
    gauss_dec
calibrator_directories :
    calibration

noise_threshold : 5

pixel_size : 0.8
        """
    with open(filename, 'w') as fobj:
        print(string, file=fobj)
    return string

def sim_map(obsdir_ra, obsdir_dec):
    simulate_map(count_map=gauss_src_func,
                 length_ra=50.,
                 length_dec=50.,
                 outdir=(obsdir_ra, obsdir_dec), mean_ra=180,
                 mean_dec=70, speed=2.,
                 spacing=0.5, srcname='Dummy')

class TestScanSet(object):
    @classmethod
    def setup_class(klass):
        import os
        global DEBUG_MODE
        DEBUG_MODE = True

        klass.curdir = os.path.dirname(__file__)
        klass.datadir = os.path.join(klass.curdir, 'data')
        klass.obsdir_ra = os.path.join(klass.datadir, 'sim', 'gauss_ra')
        klass.obsdir_dec = os.path.join(klass.datadir, 'sim', 'gauss_dec')
        klass.config_file = \
            os.path.abspath(os.path.join(klass.datadir, 'sim',
                                         'test_config_sim.ini'))
        klass.caldir = os.path.join(klass.datadir, 'sim', 'calibration')
        # First off, simulate a beamed observation  -------

        if not os.path.exists(klass.config_file):
            print('Setting up simulated data.')
            sim_config_file(klass.config_file)

        if (not os.path.exists(klass.obsdir_ra)) or \
                (not os.path.exists(klass.obsdir_dec)):
            mkdir_p(klass.obsdir_ra)
            mkdir_p(klass.obsdir_dec)
            print('Fake map: Point-like (but Gaussian beam shape), 0.5 Jy.')
            sim_map(klass.obsdir_ra, klass.obsdir_dec)

        caltable = CalibratorTable()
        caltable.from_scans(glob.glob(os.path.join(klass.caldir,
                                                   '*.fits')), debug=True)

        caltable.update()
        klass.calfile = os.path.join(klass.datadir, 'calibrators.hdf5')
        caltable.write(klass.calfile, overwrite=True)

        klass.config = read_config(klass.config_file)
        if os.path.exists('test.hdf5'):
            klass.scanset = ScanSet('test.hdf5')
        else:
            klass.scanset = ScanSet(klass.config_file, norefilt=False,
                                    debug=True)
            klass.scanset.write('test.hdf5', overwrite=True)

        plt.ioff()

    def test_0_prepare(self):
        pass

    def test_use_command_line(self):
        main_imager(('test.hdf5 -u Jy/beam '.format(self.config_file) +
                     '--calibrate {}'.format(self.calfile) +
                     ' -o bubu.hdf5 --debug').split(' '))

    def test_1_meta_saved_and_loaded_correctly(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)
        for k in scanset.meta.keys():
            assert np.all(scanset.meta[k] == self.scanset.meta[k])
        assert sorted(scanset.meta.keys()) == sorted(self.scanset.meta.keys())
        assert scanset.scan_list == self.scanset.scan_list

    def test_2_rough_image(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        images = scanset.calculate_images()

        img = images['Ch0']

        fig = plt.figure('img')
        plt.imshow(img, origin='lower')
        plt.colorbar()
        plt.savefig('img.png')
        plt.close(fig)

    def test_3_rough_image_altaz(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        images = scanset.calculate_images(altaz=True)

        img = images['Ch0']

        fig = plt.figure('img_altaz')
        plt.imshow(img, origin='lower')
        plt.colorbar()
        plt.savefig('img_altaz.png')
        scanset.save_ds9_images(save_sdev=True, altaz=True)
        plt.close(fig)

    def test_4_image_stdev(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        images = scanset.calculate_images()

        img = images['Ch0-Sdev']

        fig = plt.figure('log(img-Sdev)')
        plt.imshow(np.log10(img), origin='lower')
        plt.colorbar()
        plt.ioff()
        plt.savefig('img_sdev.png')
        plt.close(fig)

    def test_5_image_scrunch(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        images = scanset.calculate_images(scrunch=True)

        img = images['Ch0']

        fig = plt.figure('img - scrunched')
        plt.imshow(img, origin='lower')
        plt.colorbar()
        img = images['Ch0-Sdev']
        plt.savefig('img_scrunch.png')
        plt.close(fig)

        fig = plt.figure('img - scrunched - sdev')
        plt.imshow(img, origin='lower')
        plt.colorbar()
        plt.ioff()
        plt.savefig('img_scrunch_sdev.png')
        plt.close(fig)

    def test_6a_calibrate_image_pixel(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        images = scanset.calculate_images(calibration=self.calfile,
                                          map_unit="Jy/pixel")

        img = images['Ch0']
        center = img.shape[0] // 2, img.shape[1] // 2
        shortest_side = np.min(img.shape)
        X, Y = np.meshgrid(np.arange(img.shape[1]), np.arange(img.shape[0]))
        good = (X-center[1])**2 + (Y-center[0])**2 <= (shortest_side//4)**2
        assert np.allclose(np.sum(images['Ch0'][good]), 0.5, 0.05)

    def test_6b_calibrate_image_beam(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        images = scanset.calculate_images(calibration=self.calfile,
                                          map_unit="Jy/beam")

        assert np.allclose(np.max(images['Ch0']), 0.5, atol=0.05)

    def test_6c_calibrate_image_sr(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        images = scanset.calculate_images(calibration=self.calfile,
                                          map_unit="Jy/sr")

        images_pix = scanset.calculate_images(calibration=self.calfile,
                                              map_unit="Jy/pixel")

        pixel_area = scanset.meta['pixel_size'] ** 2
        assert np.allclose(images['Ch0'],
                           images_pix['Ch0'] / pixel_area.to(u.sr).value,
                           rtol=0.05)

    def test_7_ds9_image(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.save_ds9_images(save_sdev=True)

    def test_8_global_fit_image(self):
        '''Test image production.'''

        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)
        images = scanset.calculate_images()
        nx, ny = images['Ch0'].shape
        excluded = [[nx//2, ny//2, nx//4]]
        scanset.fit_full_images(excluded=excluded, chans='Ch0')
        os.path.exists("out_iter_Ch0_002.txt")
        scanset.fit_full_images(excluded=excluded, chans='Ch1')
        os.path.exists("out_iter_Ch1_000.txt")
        display_intermediate(scanset, chan="Ch0", excluded=excluded,
                             parfile="out_iter_Ch0_002.txt")
        os.path.exists("out_iter_Ch1_002.png")

    def test_9a_find_scan_through_pixel(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        _, _, _, _, _, _, _, coord = \
            scanset.find_scans_through_pixel(62, 62, test=True)

        dec_scan = os.path.join(self.obsdir_dec, 'Dec99.fits')
        assert dec_scan in coord
        assert coord[dec_scan] == 'dec'
        ra_scan = os.path.join(self.obsdir_ra, 'Ra100.fits')
        assert ra_scan in coord
        assert coord[ra_scan] == 'ra'

    def test_9b_find_scan_through_pixel(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        _, _, _, _, _, _, _, coord = \
            scanset.find_scans_through_pixel(62, 0, test=True)

        dec_scan = os.path.join(self.obsdir_dec, 'Dec99.fits')
        assert dec_scan in coord
        assert coord[dec_scan] == 'dec'
        ra_scan = os.path.join(self.obsdir_ra, 'Ra0.fits')
        assert ra_scan in coord
        assert coord[ra_scan] == 'ra'

    def test_9c_find_scan_through_invalid_pixel(self):
        scanset = ScanSet('test.hdf5',
                          config_file=self.config_file)

        scanset.calculate_images()
        _, _, _, _, _, _, _, coord = \
            scanset.find_scans_through_pixel(62, -2, test=True)
        assert coord == {}
        _, _, _, _, _, _, _, coord = \
            scanset.find_scans_through_pixel(62, 64, test=True)
        assert coord == {}

    @classmethod
    def teardown_class(klass):
        """Clean up the mess."""
        os.unlink('img.png')
        os.unlink('img_altaz.png')
        os.unlink('img_scrunch.png')
        os.unlink('delta_altaz.png')
        os.unlink('altaz_with_src.png')
        os.unlink('img_sdev.png')
        os.unlink('test.hdf5')
        os.unlink('test_scan_list.txt')
        os.unlink('img_scrunch_sdev.png')
        os.unlink('bubu.hdf5')
        for d in klass.config['list_of_directories']:
            hfiles = \
                glob.glob(os.path.join(klass.config['datadir'], d, '*.hdf5'))
            for h in hfiles:
                os.unlink(h)
        out_iter_files = glob.glob('out_iter_*')
        for o in out_iter_files:
            os.unlink(o)
        out_fits_files = glob.glob(os.path.join(klass.config['datadir'],
                                                'test_config*.fits'))
        out_hdf5_files = glob.glob(os.path.join(klass.config['datadir'], 'sim',
                                                '*/', '*.hdf5'))

        for o in out_fits_files + out_hdf5_files:
            os.unlink(o)
<|endoftext|>