<commit_before># TODO: Handle timestamp generation for modinput and set sample.timestamp properly for timestamp replacement
import datetime
import json
import os
import pprint
import random
import re
import time
import urllib.request
import urllib.parse
import urllib.error
import uuid

from splunk_eventgen.lib.timeparser import timeDelta2secs
from splunk_eventgen.lib.logging_config import logger

class Token(object):
    """Contains data and methods for replacing a token in a given sample"""
    token = None
    replacementType = None
    replacement = None
    sample = None
    mvhash = {}

    _replaytd = None
    _lastts = None
    _tokenfile = None
    _tokents = None
    _earliestTime = None
    _latestTime = None
    _replacementFile = None
    _replacementColumn = None
    _integerMatch = None
    _floatMatch = None
    _hexMatch = None
    _stringMatch = None
    _listMatch = None
    _tokenfilecounter = 0

    def __init__(self, sample=None):
        self._earliestTime = (None, None)
        self._latestTime = (None, None)

    def __str__(self):
        """Only used for debugging, outputs a pretty printed representation of this token"""
        # Eliminate recursive going back to parent
        temp = dict([(key, value) for (key, value) in self.__dict__.items() if key != 'sample'])
        return pprint.pformat(temp)

    def __repr__(self):
        return self.__str__()

    def _match(self, event):
        """Executes regular expression match and returns the re.Match object"""
        return re.match(self.token, event)

    def _search(self, event):
        """Executes regular expression search and returns the re.Match object"""
        return re.search(self.token, event)

    def _finditer(self, event):
        """Executes regular expression finditer and returns the re.Match object"""
        return re.finditer(self.token, event)

    def _findall(self, event):
        """Executes regular expression finditer and returns the re.Match object"""
        return re.findall(self.token, event)

    def replace(self, event, et=None, lt=None, s=None, pivot_timestamp=None):
        """Replaces all instances of this token in provided event and returns event"""
        offset = 0
        tokenMatch = list(self._finditer(event))

        if len(tokenMatch) > 0:
            replacement = self._getReplacement(event[tokenMatch[0].start(0):tokenMatch[0].end(0)], et, lt, s,
                                               pivot_timestamp=pivot_timestamp)
            if replacement is not None or self.replacementType == 'replaytimestamp':
                # logger.debug("Replacement: '%s'" % replacement)
                # Iterate matches
                for match in tokenMatch:
                    # logger.debug("Match: %s" % (match))
                    try:
                        matchStart = match.start(1) + offset
                        matchEnd = match.end(1) + offset
                        startEvent = event[:matchStart]
                        endEvent = event[matchEnd:]
                        # In order to not break legacy which might replace the same timestamp
                        # with the same value in multiple matches, here we'll include
                        # ones that need to be replaced for every match
                        if self.replacementType == 'replaytimestamp':
                            replacement = lt.strftime(self.replacement)
                        offset += len(replacement) - len(match.group(1))
                    except:
                        matchStart = match.start(0) + offset
                        matchEnd = match.end(0) + offset
                        startEvent = event[:matchStart]
                        endEvent = event[matchEnd:]
                        # In order to not break legacy which might replace the same timestamp
                        # with the same value in multiple matches, here we'll include
                        # ones that need to be replaced for every match
                        if self.replacementType == 'replaytimestamp':
                            replacement = lt.strftime(self.replacement)
                        offset += len(replacement) - len(match.group(0))
                    # logger.debug("matchStart %d matchEnd %d offset %d" % (matchStart, matchEnd, offset))
                    event = startEvent + replacement + endEvent

                # Reset replay internal variables for this token
                self._replaytd = None
                self._lastts = None
        return event

    def _getReplacement(self, old=None, earliestTime=None, latestTime=None, s=None, pivot_timestamp=None):
        if self.replacementType == 'static':
            return self.replacement
        # This logic is done in replay.py
        elif self.replacementType == 'replaytimestamp':
            pass
        elif self.replacementType == 'timestamp':
            if s.earliest and s.latest:
                if earliestTime and latestTime:
                    if latestTime >= earliestTime:
                        if pivot_timestamp:
                            replacementTime = pivot_timestamp
                        elif s.timestamp is None:
                            minDelta = 0

                            # Compute timeDelta as total_seconds
                            td = latestTime - earliestTime
                            if not type(td) == float:
                                maxDelta = timeDelta2secs(td)
                            else:
                                maxDelta = td

                            # Get random timeDelta
                            randomDelta = datetime.timedelta(
                                seconds=random.randint(minDelta, maxDelta), microseconds=random.randint(
                                    0, latestTime.microsecond if latestTime.microsecond > 0 else 999999))

                            # Compute replacmentTime
                            replacementTime = latestTime - randomDelta
                            s.timestamp = replacementTime
                        else:
                            replacementTime = s.timestamp

                        replacement = self.replacement.replace(
                            '%s',
                            str(round(time.mktime(replacementTime.timetuple()))).rstrip('0').rstrip('.'))
                        replacementTime = replacementTime.strftime(replacement)
                        # replacementTime == replacement for invalid strptime specifiers
                        if replacementTime != self.replacement.replace('%', ''):
                            return replacementTime
                        else:
                            logger.error(
                                "Invalid strptime specifier '%s' detected; will not replace" % (self.replacement))
                            return old
                    # earliestTime/latestTime not proper
                    else:
                        logger.error(("Earliest specifier '%s', value '%s' is greater than latest specifier '%s'" +
                                           "value '%s' for sample '%s'; will not replace") %
                                          (s.earliest, earliestTime, s.latest, latestTime, s.name))
                        return old
            # earliest/latest not proper
            else:
                logger.error('Earliest or latest specifier were not set; will not replace')
                return old
        elif self.replacementType in ('random', 'rated'):
            # Validations:
            if self._integerMatch is not None:
                integerMatch = self._integerMatch
            else:
                integerRE = re.compile(r'integer\[([-]?\d+):([-]?\d+)\]', re.I)
                integerMatch = integerRE.match(self.replacement)
                self._integerMatch = integerMatch

            if self._floatMatch is not None:
                floatMatch = self._floatMatch
            else:
                floatRE = re.compile(r'float\[(-?\d+|-?\d+\.(\d+)):(-?\d+|-?\d+\.(\d+))\]', re.I)
                floatMatch = floatRE.match(self.replacement)
                self._floatMatch = floatMatch

            if self._stringMatch is not None:
                stringMatch = self._stringMatch
            else:
                stringRE = re.compile(r'string\((\d+)\)', re.I)
                stringMatch = stringRE.match(self.replacement)
                self._stringMatch = stringMatch

            if self._hexMatch is not None:
                hexMatch = self._hexMatch
            else:
                hexRE = re.compile(r'hex\((\d+)\)', re.I)
                hexMatch = hexRE.match(self.replacement)
                self._hexMatch = hexMatch

            if self._listMatch is not None:
                listMatch = self._listMatch
            else:
                listRE = re.compile(r'list(\[[^\]]+\])', re.I)
                listMatch = listRE.match(self.replacement)
                self._listMatch = listMatch

            # Valid replacements: ipv4 | ipv6 | integer[<start>:<end>] | string(<i>)
            if self.replacement.lower() == 'ipv4':
                x = 0
                replacement = ''

                while x < 4:
                    replacement += str(random.randint(0, 255)) + '.'
                    x += 1

                replacement = replacement.strip('.')
                return replacement
            elif self.replacement.lower() == 'ipv6':
                x = 0
                replacement = ''

                while x < 8:
                    replacement += hex(random.randint(0, 65535))[2:] + ':'
                    x += 1

                replacement = replacement.strip(':')
                return replacement
            elif self.replacement.lower() == 'mac':
                x = 0
                replacement = ''

                # Give me 6 blocks of 2 hex
                while x < 6:
                    y = 0
                    while y < 2:
                        replacement += hex(random.randint(0, 15))[2:]
                        y += 1
                    replacement += ':'
                    x += 1

                replacement = replacement.strip(':')
                return replacement
            elif self.replacement.lower() == 'guid':
                return str(uuid.uuid4())
            elif integerMatch:
                startInt = int(integerMatch.group(1))
                endInt = int(integerMatch.group(2))

                if endInt >= startInt:
                    replacementInt = random.randint(startInt, endInt)
                    if self.replacementType == 'rated':
                        rateFactor = 1.0
                        if type(s.hourOfDayRate) == dict:
                            try:
                                rateFactor *= s.hourOfDayRate[str(s.now())]
                            except KeyError:
                                import traceback
                                stack = traceback.format_exc()
                                logger.error("Hour of day rate failed for token %s.  Stacktrace %s" % stack)
                        if type(s.dayOfWeekRate) == dict:
                            try:
                                weekday = datetime.date.weekday(s.now())
                                if weekday == 6:
                                    weekday = 0
                                else:
                                    weekday += 1
                                rateFactor *= s.dayOfWeekRate[str(weekday)]
                            except KeyError:
                                import traceback
                                stack = traceback.format_exc()
                                logger.error("Day of week rate failed.  Stacktrace %s" % stack)
                        replacementInt = int(round(replacementInt * rateFactor, 0))
                    replacement = str(replacementInt)
                    return replacement
                else:
                    logger.error(
                        "Start integer %s greater than end integer %s; will not replace" % (startInt, endInt))
                    return old
            elif floatMatch:
                try:
                    startFloat = float(floatMatch.group(1))
                    endFloat = float(floatMatch.group(3))

                    significance = 0
                    if floatMatch.group(2) is not None:
                        significance = len(floatMatch.group(2))

                    if endFloat >= startFloat:
                        floatret = round(random.uniform(startFloat, endFloat), significance)
                        if self.replacementType == 'rated':
                            rateFactor = 1.0
                            now = s.now()
                            if type(s.hourOfDayRate) == dict:
                                try:
                                    rateFactor *= s.hourOfDayRate[str(now.hour)]
                                except KeyError:
                                    import traceback
                                    stack = traceback.format_exc()
                                    logger.error("Hour of day rate failed for token %s.  Stacktrace %s" % stack)
                            if type(s.dayOfWeekRate) == dict:
                                try:
                                    weekday = datetime.date.weekday(now)
                                    if weekday == 6:
                                        weekday = 0
                                    else:
                                        weekday += 1
                                    rateFactor *= s.dayOfWeekRate[str(weekday)]
                                except KeyError:
                                    import traceback
                                    stack = traceback.format_exc()
                                    logger.error("Day of week rate failed.  Stacktrace %s" % stack)
                            floatret = round(floatret * rateFactor, significance)
                        floatret = str(floatret)
                        return floatret
                    else:
                        logger.error(
                            "Start float %s greater than end float %s; will not replace" % (startFloat, endFloat))
                        return old
                except ValueError:
                    logger.error("Could not parse float[%s:%s]" % (floatMatch.group(1), floatMatch.group(4)))
                    return old
            elif stringMatch:
                strLength = int(stringMatch.group(1))
                if strLength == 0:
                    return ''
                elif strLength > 0:
                    replacement = ''
                    while len(replacement) < strLength:
                        # Generate a random ASCII between dec 33->126
                        replacement += chr(random.randint(33, 126))
                        # Practice safe strings
                        replacement = re.sub('%[0-9a-fA-F]+', '', urllib.parse.quote(replacement))

                    return replacement
                else:
                    logger.error(
                        "Length specifier %s for string replacement must be greater than 0; will not replace" %
                        (strLength))
                    return old
            elif hexMatch:
                strLength = int(hexMatch.group(1))

                replacement = ''
                hexList = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F']
                while len(replacement) < strLength:
                    replacement += hexList[random.randint(0, 15)]

                return replacement
            elif listMatch:
                try:
                    value = json.loads(listMatch.group(1))
                except:
                    logger.error("Could not parse json for '%s' in sample '%s'" % (listMatch.group(1), s.name))
                    return old
                return random.SystemRandom().choice(value)

            else:
                logger.error("Unknown replacement value '%s' for replacementType '%s'; will not replace" %
                                  (self.replacement, self.replacementType))
                return old
        elif self.replacementType in ('file', 'mvfile', 'seqfile'):
            if self._replacementFile is not None:
                replacementFile = self._replacementFile
                replacementColumn = self._replacementColumn
            else:
                try:
                    paths = self.replacement.split(':')
                    if (len(paths) == 1):
                        replacementColumn = 0
                    else:
                        try:  # When it's not a mvfile, there's no number on the end:
                            replacementColumn = int(paths[-1])
                        except (ValueError):
                            replacementColumn = 0
                    if (replacementColumn > 0):
                        # This supports having a drive-letter colon
                        replacementFile = s.pathParser(":".join(paths[0:-1]))
                    else:
                        replacementFile = s.pathParser(self.replacement)
                except ValueError:
                    logger.error(
                        "Replacement string '%s' improperly formatted. Should be /path/to/file or /path/to/file:column"
                        % self.replacement)
                    return old
                self._replacementFile = replacementFile
                self._replacementColumn = replacementColumn

            # If we've seen this file before, simply return already read results
            # This applies only if we're looking at a multivalue file and we want to
            # return the same random pick on every iteration
            if replacementColumn > 0 and replacementFile in self.mvhash:
                if replacementColumn > len(self.mvhash[replacementFile]):
                    logger.error("Index for column '%s' in replacement file '%s' is out of bounds" %
                                      (replacementColumn, replacementFile))
                    return old
                else:
                    # logger.debug("Returning mvhash: %s" % self.mvhash[replacementFile][replacementColumn-1])
                    return self.mvhash[replacementFile][replacementColumn - 1]
            else:
                # Adding caching of the token file to avoid reading it every iteration
                if self._tokenfile is not None:
                    replacementLines = self._tokenfile
                # Otherwise, lets read the file and build our cached results, pick a result and return it
                else:
                    # logger.debug("replacementFile: %s replacementColumn: %s" %
                    #                   (replacementFile, replacementColumn))
                    replacementFile = os.path.abspath(replacementFile)
                    logger.debug("Normalized replacement file %s" % replacementFile)
                    if os.path.exists(replacementFile) and os.path.isfile(replacementFile):
                        replacementFH = open(replacementFile, 'rU')
                        replacementLines = replacementFH.readlines()
                        replacementFH.close()

                        if len(replacementLines) == 0:
                            logger.error("Replacement file '%s' is empty; will not replace" % (replacementFile))
                            return old
                        else:
                            self._tokenfile = replacementLines
                    else:
                        logger.error("File '%s' does not exist" % (replacementFile))
                        return old
                if self.replacementType == 'seqfile':
                    # pick value one by one from replacement file
                    replacement = replacementLines[self._tokenfilecounter % len(replacementLines)].strip()
                    self._tokenfilecounter += 1
                else:
                    # pick value randomly from replacement file
                    replacement = replacementLines[random.randint(0, len(replacementLines) - 1)].strip()

                if replacementColumn > 0:
                    self.mvhash[replacementFile] = replacement.split(',')

                    if replacementColumn > len(self.mvhash[replacementFile]):
                        logger.error("Index for column '%s' in replacement file '%s' is out of bounds" %
                                          (replacementColumn, replacementFile))
                        return old
                    else:
                        return self.mvhash[replacementFile][replacementColumn - 1]
                else:
                    return replacement
        elif self.replacementType == 'integerid':
            temp = self.replacement
            self.replacement = str(int(self.replacement) + 1)
            return temp

        else:
            logger.error("Unknown replacementType '%s'; will not replace" % self.replacementType)
            return old
<commit_msg>Fix high system load issue (#345)<commit_after># TODO: Handle timestamp generation for modinput and set sample.timestamp properly for timestamp replacement
import datetime
import json
import os
import pprint
import random
import re
import time
import urllib.request
import urllib.parse
import urllib.error
import uuid

from splunk_eventgen.lib.timeparser import timeDelta2secs
from splunk_eventgen.lib.logging_config import logger

class Token(object):
    """Contains data and methods for replacing a token in a given sample"""
    token = None
    replacementType = None
    replacement = None
    sample = None
    mvhash = {}

    _replaytd = None
    _lastts = None
    _tokenfile = None
    _tokents = None
    _earliestTime = None
    _latestTime = None
    _replacementFile = None
    _replacementColumn = None
    _integerMatch = None
    _floatMatch = None
    _hexMatch = None
    _stringMatch = None
    _listMatch = None
    _tokenfilecounter = 0

    def __init__(self, sample=None):
        self._earliestTime = (None, None)
        self._latestTime = (None, None)

    def __str__(self):
        """Only used for debugging, outputs a pretty printed representation of this token"""
        # Eliminate recursive going back to parent
        temp = dict([(key, value) for (key, value) in self.__dict__.items() if key != 'sample'])
        return pprint.pformat(temp)

    def __repr__(self):
        return self.__str__()

    def _match(self, event):
        """Executes regular expression match and returns the re.Match object"""
        return re.match(self.token, event)

    def _search(self, event):
        """Executes regular expression search and returns the re.Match object"""
        return re.search(self.token, event)

    def _finditer(self, event):
        """Executes regular expression finditer and returns the re.Match object"""
        return re.finditer(self.token, event)

    def _findall(self, event):
        """Executes regular expression finditer and returns the re.Match object"""
        return re.findall(self.token, event)

    def replace(self, event, et=None, lt=None, s=None, pivot_timestamp=None):
        """Replaces all instances of this token in provided event and returns event"""
        offset = 0
        tokenMatch = list(self._finditer(event))

        if len(tokenMatch) > 0:
            replacement = self._getReplacement(event[tokenMatch[0].start(0):tokenMatch[0].end(0)], et, lt, s,
                                               pivot_timestamp=pivot_timestamp)
            if replacement is not None or self.replacementType == 'replaytimestamp':
                # logger.debug("Replacement: '%s'" % replacement)
                # Iterate matches
                for match in tokenMatch:
                    # logger.debug("Match: %s" % (match))
                    try:
                        matchStart = match.start(1) + offset
                        matchEnd = match.end(1) + offset
                        startEvent = event[:matchStart]
                        endEvent = event[matchEnd:]
                        # In order to not break legacy which might replace the same timestamp
                        # with the same value in multiple matches, here we'll include
                        # ones that need to be replaced for every match
                        if self.replacementType == 'replaytimestamp':
                            replacement = lt.strftime(self.replacement)
                        offset += len(replacement) - len(match.group(1))
                    except:
                        matchStart = match.start(0) + offset
                        matchEnd = match.end(0) + offset
                        startEvent = event[:matchStart]
                        endEvent = event[matchEnd:]
                        # In order to not break legacy which might replace the same timestamp
                        # with the same value in multiple matches, here we'll include
                        # ones that need to be replaced for every match
                        if self.replacementType == 'replaytimestamp':
                            replacement = lt.strftime(self.replacement)
                        offset += len(replacement) - len(match.group(0))
                    # logger.debug("matchStart %d matchEnd %d offset %d" % (matchStart, matchEnd, offset))
                    event = startEvent + replacement + endEvent

                # Reset replay internal variables for this token
                self._replaytd = None
                self._lastts = None
        return event

    def _getReplacement(self, old=None, earliestTime=None, latestTime=None, s=None, pivot_timestamp=None):
        if self.replacementType == 'static':
            return self.replacement
        # This logic is done in replay.py
        elif self.replacementType == 'replaytimestamp':
            pass
        elif self.replacementType == 'timestamp':
            if s.earliest and s.latest:
                if earliestTime and latestTime:
                    if latestTime >= earliestTime:
                        if pivot_timestamp:
                            replacementTime = pivot_timestamp
                        elif s.timestamp is None:
                            minDelta = 0

                            # Compute timeDelta as total_seconds
                            td = latestTime - earliestTime
                            if not type(td) == float:
                                maxDelta = timeDelta2secs(td)
                            else:
                                maxDelta = td

                            # Get random timeDelta
                            randomDelta = datetime.timedelta(
                                seconds=random.randint(minDelta, maxDelta), microseconds=random.randint(
                                    0, latestTime.microsecond if latestTime.microsecond > 0 else 999999))

                            # Compute replacmentTime
                            replacementTime = latestTime - randomDelta
                            s.timestamp = replacementTime
                        else:
                            replacementTime = s.timestamp

                        replacement = self.replacement.replace(
                            '%s',
                            str(round(time.mktime(replacementTime.timetuple()))).rstrip('0').rstrip('.'))
                        replacementTime = replacementTime.strftime(replacement)
                        # replacementTime == replacement for invalid strptime specifiers
                        if replacementTime != self.replacement.replace('%', ''):
                            return replacementTime
                        else:
                            logger.error(
                                "Invalid strptime specifier '%s' detected; will not replace" % (self.replacement))
                            return old
                    # earliestTime/latestTime not proper
                    else:
                        logger.error(("Earliest specifier '%s', value '%s' is greater than latest specifier '%s'" +
                                           "value '%s' for sample '%s'; will not replace") %
                                          (s.earliest, earliestTime, s.latest, latestTime, s.name))
                        return old
            # earliest/latest not proper
            else:
                logger.error('Earliest or latest specifier were not set; will not replace')
                return old
        elif self.replacementType in ('random', 'rated'):
            # Validations:
            if self._integerMatch is not None:
                integerMatch = self._integerMatch
            else:
                integerRE = re.compile(r'integer\[([-]?\d+):([-]?\d+)\]', re.I)
                integerMatch = integerRE.match(self.replacement)
                self._integerMatch = integerMatch

            if self._floatMatch is not None:
                floatMatch = self._floatMatch
            else:
                floatRE = re.compile(r'float\[(-?\d+|-?\d+\.(\d+)):(-?\d+|-?\d+\.(\d+))\]', re.I)
                floatMatch = floatRE.match(self.replacement)
                self._floatMatch = floatMatch

            if self._stringMatch is not None:
                stringMatch = self._stringMatch
            else:
                stringRE = re.compile(r'string\((\d+)\)', re.I)
                stringMatch = stringRE.match(self.replacement)
                self._stringMatch = stringMatch

            if self._hexMatch is not None:
                hexMatch = self._hexMatch
            else:
                hexRE = re.compile(r'hex\((\d+)\)', re.I)
                hexMatch = hexRE.match(self.replacement)
                self._hexMatch = hexMatch

            if self._listMatch is not None:
                listMatch = self._listMatch
            else:
                listRE = re.compile(r'list(\[[^\]]+\])', re.I)
                listMatch = listRE.match(self.replacement)
                self._listMatch = listMatch

            # Valid replacements: ipv4 | ipv6 | integer[<start>:<end>] | string(<i>)
            if self.replacement.lower() == 'ipv4':
                x = 0
                replacement = ''

                while x < 4:
                    replacement += str(random.randint(0, 255)) + '.'
                    x += 1

                replacement = replacement.strip('.')
                return replacement
            elif self.replacement.lower() == 'ipv6':
                x = 0
                replacement = ''

                while x < 8:
                    replacement += hex(random.randint(0, 65535))[2:] + ':'
                    x += 1

                replacement = replacement.strip(':')
                return replacement
            elif self.replacement.lower() == 'mac':
                x = 0
                replacement = ''

                # Give me 6 blocks of 2 hex
                while x < 6:
                    y = 0
                    while y < 2:
                        replacement += hex(random.randint(0, 15))[2:]
                        y += 1
                    replacement += ':'
                    x += 1

                replacement = replacement.strip(':')
                return replacement
            elif self.replacement.lower() == 'guid':
                return str(uuid.uuid4())
            elif integerMatch:
                startInt = int(integerMatch.group(1))
                endInt = int(integerMatch.group(2))

                if endInt >= startInt:
                    replacementInt = random.randint(startInt, endInt)
                    if self.replacementType == 'rated':
                        rateFactor = 1.0
                        if type(s.hourOfDayRate) == dict:
                            try:
                                rateFactor *= s.hourOfDayRate[str(s.now())]
                            except KeyError:
                                import traceback
                                stack = traceback.format_exc()
                                logger.error("Hour of day rate failed for token %s.  Stacktrace %s" % stack)
                        if type(s.dayOfWeekRate) == dict:
                            try:
                                weekday = datetime.date.weekday(s.now())
                                if weekday == 6:
                                    weekday = 0
                                else:
                                    weekday += 1
                                rateFactor *= s.dayOfWeekRate[str(weekday)]
                            except KeyError:
                                import traceback
                                stack = traceback.format_exc()
                                logger.error("Day of week rate failed.  Stacktrace %s" % stack)
                        replacementInt = int(round(replacementInt * rateFactor, 0))
                    replacement = str(replacementInt)
                    return replacement
                else:
                    logger.error(
                        "Start integer %s greater than end integer %s; will not replace" % (startInt, endInt))
                    return old
            elif floatMatch:
                try:
                    startFloat = float(floatMatch.group(1))
                    endFloat = float(floatMatch.group(3))

                    significance = 0
                    if floatMatch.group(2) is not None:
                        significance = len(floatMatch.group(2))

                    if endFloat >= startFloat:
                        floatret = round(random.uniform(startFloat, endFloat), significance)
                        if self.replacementType == 'rated':
                            rateFactor = 1.0
                            now = s.now()
                            if type(s.hourOfDayRate) == dict:
                                try:
                                    rateFactor *= s.hourOfDayRate[str(now.hour)]
                                except KeyError:
                                    import traceback
                                    stack = traceback.format_exc()
                                    logger.error("Hour of day rate failed for token %s.  Stacktrace %s" % stack)
                            if type(s.dayOfWeekRate) == dict:
                                try:
                                    weekday = datetime.date.weekday(now)
                                    if weekday == 6:
                                        weekday = 0
                                    else:
                                        weekday += 1
                                    rateFactor *= s.dayOfWeekRate[str(weekday)]
                                except KeyError:
                                    import traceback
                                    stack = traceback.format_exc()
                                    logger.error("Day of week rate failed.  Stacktrace %s" % stack)
                            floatret = round(floatret * rateFactor, significance)
                        floatret = str(floatret)
                        return floatret
                    else:
                        logger.error(
                            "Start float %s greater than end float %s; will not replace" % (startFloat, endFloat))
                        return old
                except ValueError:
                    logger.error("Could not parse float[%s:%s]" % (floatMatch.group(1), floatMatch.group(4)))
                    return old
            elif stringMatch:
                strLength = int(stringMatch.group(1))
                if strLength == 0:
                    return ''
                elif strLength > 0:
                    replacement = ''
                    while len(replacement) < strLength:
                        # Generate a random ASCII between dec 33->126
                        replacement += chr(random.randint(33, 126))
                        # Practice safe strings
                        replacement = re.sub('%[0-9a-fA-F]+', '', urllib.parse.quote(replacement))

                    return replacement
                else:
                    logger.error(
                        "Length specifier %s for string replacement must be greater than 0; will not replace" %
                        (strLength))
                    return old
            elif hexMatch:
                strLength = int(hexMatch.group(1))

                replacement = ''
                hexList = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F']
                while len(replacement) < strLength:
                    replacement += hexList[random.randint(0, 15)]

                return replacement
            elif listMatch:
                try:
                    value = json.loads(listMatch.group(1))
                except:
                    logger.error("Could not parse json for '%s' in sample '%s'" % (listMatch.group(1), s.name))
                    return old
                return random.choice(value)

            else:
                logger.error("Unknown replacement value '%s' for replacementType '%s'; will not replace" %
                                  (self.replacement, self.replacementType))
                return old
        elif self.replacementType in ('file', 'mvfile', 'seqfile'):
            if self._replacementFile is not None:
                replacementFile = self._replacementFile
                replacementColumn = self._replacementColumn
            else:
                try:
                    paths = self.replacement.split(':')
                    if (len(paths) == 1):
                        replacementColumn = 0
                    else:
                        try:  # When it's not a mvfile, there's no number on the end:
                            replacementColumn = int(paths[-1])
                        except (ValueError):
                            replacementColumn = 0
                    if (replacementColumn > 0):
                        # This supports having a drive-letter colon
                        replacementFile = s.pathParser(":".join(paths[0:-1]))
                    else:
                        replacementFile = s.pathParser(self.replacement)
                except ValueError:
                    logger.error(
                        "Replacement string '%s' improperly formatted. Should be /path/to/file or /path/to/file:column"
                        % self.replacement)
                    return old
                self._replacementFile = replacementFile
                self._replacementColumn = replacementColumn

            # If we've seen this file before, simply return already read results
            # This applies only if we're looking at a multivalue file and we want to
            # return the same random pick on every iteration
            if replacementColumn > 0 and replacementFile in self.mvhash:
                if replacementColumn > len(self.mvhash[replacementFile]):
                    logger.error("Index for column '%s' in replacement file '%s' is out of bounds" %
                                      (replacementColumn, replacementFile))
                    return old
                else:
                    # logger.debug("Returning mvhash: %s" % self.mvhash[replacementFile][replacementColumn-1])
                    return self.mvhash[replacementFile][replacementColumn - 1]
            else:
                # Adding caching of the token file to avoid reading it every iteration
                if self._tokenfile is not None:
                    replacementLines = self._tokenfile
                # Otherwise, lets read the file and build our cached results, pick a result and return it
                else:
                    # logger.debug("replacementFile: %s replacementColumn: %s" %
                    #                   (replacementFile, replacementColumn))
                    replacementFile = os.path.abspath(replacementFile)
                    logger.debug("Normalized replacement file %s" % replacementFile)
                    if os.path.exists(replacementFile) and os.path.isfile(replacementFile):
                        replacementFH = open(replacementFile, 'rU')
                        replacementLines = replacementFH.readlines()
                        replacementFH.close()

                        if len(replacementLines) == 0:
                            logger.error("Replacement file '%s' is empty; will not replace" % (replacementFile))
                            return old
                        else:
                            self._tokenfile = replacementLines
                    else:
                        logger.error("File '%s' does not exist" % (replacementFile))
                        return old
                if self.replacementType == 'seqfile':
                    # pick value one by one from replacement file
                    replacement = replacementLines[self._tokenfilecounter % len(replacementLines)].strip()
                    self._tokenfilecounter += 1
                else:
                    # pick value randomly from replacement file
                    replacement = replacementLines[random.randint(0, len(replacementLines) - 1)].strip()

                if replacementColumn > 0:
                    self.mvhash[replacementFile] = replacement.split(',')

                    if replacementColumn > len(self.mvhash[replacementFile]):
                        logger.error("Index for column '%s' in replacement file '%s' is out of bounds" %
                                          (replacementColumn, replacementFile))
                        return old
                    else:
                        return self.mvhash[replacementFile][replacementColumn - 1]
                else:
                    return replacement
        elif self.replacementType == 'integerid':
            temp = self.replacement
            self.replacement = str(int(self.replacement) + 1)
            return temp

        else:
            logger.error("Unknown replacementType '%s'; will not replace" % self.replacementType)
            return old
<|endoftext|>


<commit_before>import pytest
import numpy as np

from astropy import units as u
from astropy import convolution
from astropy.wcs import WCS
from astropy import wcs
from astropy.io import fits

from radio_beam import beam, Beam

from .. import SpectralCube
from ..utils import WCSCelestialError
from .test_spectral_cube import cube_and_raw
from .test_projection import load_projection
from . import path

def test_convolution():
    cube, data = cube_and_raw('255_delta.fits')

    # 1" convolved with 1.5" -> 1.8027....
    target_beam = Beam(1.802775637731995*u.arcsec, 1.802775637731995*u.arcsec,
                       0*u.deg)

    conv_cube = cube.convolve_to(target_beam)

    expected = convolution.Gaussian2DKernel((1.5*u.arcsec /
                                             beam.SIGMA_TO_FWHM /
                                             (5.555555555555e-4*u.deg)).decompose().value,
                                            x_size=5, y_size=5,
                                           )

    expected.normalize()

    np.testing.assert_almost_equal(expected.array,
                                   conv_cube.filled_data[0,:,:].value)

    # 2nd layer is all zeros
    assert np.all(conv_cube.filled_data[1,:,:] == 0.0)

def test_beams_convolution():
    cube, data = cube_and_raw('455_delta_beams.fits')

    # 1" convolved with 1.5" -> 1.8027....
    target_beam = Beam(1.802775637731995*u.arcsec, 1.802775637731995*u.arcsec,
                       0*u.deg)

    conv_cube = cube.convolve_to(target_beam)

    pixscale = wcs.utils.proj_plane_pixel_area(cube.wcs.celestial)**0.5*u.deg

    for ii, bm in enumerate(cube.beams):
        expected = target_beam.deconvolve(bm).as_kernel(pixscale, x_size=5,
                                                        y_size=5)
        expected.normalize()

        np.testing.assert_almost_equal(expected.array,
                                       conv_cube.filled_data[ii,:,:].value)

def test_beams_convolution_equal():
    cube, data = cube_and_raw('522_delta_beams.fits')

    # Only checking that the equal beam case is handled correctly.
    # Fake the beam in the first channel. Then ensure that the first channel
    # has NOT been convolved.
    target_beam = Beam(1.0 * u.arcsec, 1.0 * u.arcsec, 0.0 * u.deg)
    cube.beams.major[0] = target_beam.major
    cube.beams.minor[0] = target_beam.minor
    cube.beams.pa[0] = target_beam.pa

    conv_cube = cube.convolve_to(target_beam)

    np.testing.assert_almost_equal(cube.filled_data[0].value,
                                   conv_cube.filled_data[0].value)

def test_reproject():

    pytest.importorskip('reproject')

    cube, data = cube_and_raw('adv.fits')

    wcs_in = WCS(cube.header)
    wcs_out = wcs_in.deepcopy()
    wcs_out.wcs.ctype = ['GLON-SIN', 'GLAT-SIN', wcs_in.wcs.ctype[2]]
    wcs_out.wcs.crval = [134.37608, -31.939241, wcs_in.wcs.crval[2]]
    wcs_out.wcs.crpix = [2., 2., wcs_in.wcs.crpix[2]]

    header_out = cube.header
    header_out['NAXIS1'] = 4
    header_out['NAXIS2'] = 5
    header_out['NAXIS3'] = cube.shape[0]
    header_out.update(wcs_out.to_header())

    result = cube.reproject(header_out)

    assert result.shape == (cube.shape[0], 5, 4)

def test_spectral_smooth():

    cube, data = cube_and_raw('522_delta.fits')

    result = cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0), use_memmap=False)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   convolution.Gaussian1DKernel(1.0,
                                                                x_size=5).array,
                                   4)

    result = cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0), use_memmap=True)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   convolution.Gaussian1DKernel(1.0,
                                                                x_size=5).array,
                                   4)

def test_spectral_smooth_4cores():

    pytest.importorskip('joblib')

    cube, data = cube_and_raw('522_delta.fits')

    result = cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0), num_cores=4, use_memmap=True)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   convolution.Gaussian1DKernel(1.0,
                                                                x_size=5).array,
                                   4)

    # this is one way to test non-parallel mode
    result = cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0), num_cores=4, use_memmap=False)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   convolution.Gaussian1DKernel(1.0,
                                                                x_size=5).array,
                                   4)

    # num_cores = 4 is a contradiction with parallel=False, but we want to make
    # sure it does the same thing
    result = cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0), num_cores=4, parallel=False)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   convolution.Gaussian1DKernel(1.0,
                                                                x_size=5).array,
                                   4)

def test_spectral_smooth_fail():

    cube, data = cube_and_raw('522_delta_beams.fits')

    with pytest.raises(AttributeError) as exc:
        cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0))

    assert exc.value.args[0] == ("VaryingResolutionSpectralCubes can't be "
                                 "spectrally smoothed.  Convolve to a "
                                 "common resolution with `convolve_to` before "
                                 "attempting spectral smoothed.")

def test_spectral_interpolate():

    cube, data = cube_and_raw('522_delta.fits')

    orig_wcs = cube.wcs.deepcopy()

    # midpoint between each position
    sg = (cube.spectral_axis[1:] + cube.spectral_axis[:-1])/2.

    result = cube.spectral_interpolate(spectral_grid=sg)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   [0.0, 0.5, 0.5, 0.0])

    assert cube.wcs.wcs.compare(orig_wcs.wcs)

def test_spectral_interpolate_with_fillvalue():

    cube, data = cube_and_raw('522_delta.fits')

    # Step one channel out of bounds.
    sg = ((cube.spectral_axis[0]) -
          (cube.spectral_axis[1] - cube.spectral_axis[0]) *
          np.linspace(1,4,4))
    result = cube.spectral_interpolate(spectral_grid=sg,
                                       fill_value=42)
    np.testing.assert_almost_equal(result[:,0,0].value,
                                   np.ones(4)*42)

def test_spectral_interpolate_fail():

    cube, data = cube_and_raw('522_delta_beams.fits')

    with pytest.raises(AttributeError) as exc:
        cube.spectral_interpolate(5)

    assert exc.value.args[0] == ("VaryingResolutionSpectralCubes can't be "
                                 "spectrally interpolated.  Convolve to a "
                                 "common resolution with `convolve_to` before "
                                 "attempting spectral interpolation.")

def test_spectral_interpolate_with_mask():

    hdu = fits.open(path("522_delta.fits"))[0]

    # Swap the velocity axis so indiff < 0 in spectral_interpolate
    hdu.header["CDELT3"] = - hdu.header["CDELT3"]

    cube = SpectralCube.read(hdu)

    mask = np.ones(cube.shape, dtype=bool)
    mask[:2] = False

    masked_cube = cube.with_mask(mask)

    orig_wcs = cube.wcs.deepcopy()

    # midpoint between each position
    sg = (cube.spectral_axis[1:] + cube.spectral_axis[:-1])/2.

    result = masked_cube.spectral_interpolate(spectral_grid=sg[::-1])

    # The output makes CDELT3 > 0 (reversed spectral axis) so the masked
    # portion are the final 2 channels.
    np.testing.assert_almost_equal(result[:, 0, 0].value,
                                   [0.0, 0.5, np.NaN, np.NaN])

    assert cube.wcs.wcs.compare(orig_wcs.wcs)

def test_spectral_interpolate_reversed():

    cube, data = cube_and_raw('522_delta.fits')

    orig_wcs = cube.wcs.deepcopy()

    # Reverse spectral axis
    sg = cube.spectral_axis[::-1]

    result = cube.spectral_interpolate(spectral_grid=sg)

    np.testing.assert_almost_equal(sg.value, result.spectral_axis.value)

def test_convolution_2D():

    proj, hdu = load_projection("55_delta.fits")

    # 1" convolved with 1.5" -> 1.8027....
    target_beam = Beam(1.802775637731995*u.arcsec, 1.802775637731995*u.arcsec,
                       0*u.deg)

    conv_proj = proj.convolve_to(target_beam)

    expected = convolution.Gaussian2DKernel((1.5*u.arcsec /
                                             beam.SIGMA_TO_FWHM /
                                             (5.555555555555e-4*u.deg)).decompose().value,
                                            x_size=5, y_size=5,
                                           )
    expected.normalize()

    np.testing.assert_almost_equal(expected.array,
                                   conv_proj.value)
    assert conv_proj.beam == target_beam

def test_nocelestial_convolution_2D_fail():

    cube, data = cube_and_raw('255_delta.fits')

    proj = cube.moment0(axis=1)

    test_beam = Beam(1.0 * u.arcsec)

    with pytest.raises(WCSCelestialError) as exc:
        proj.convolve_to(test_beam)

    assert exc.value.args[0] == ("WCS does not contain two spatial axes.")

def test_reproject_2D():

    pytest.importorskip('reproject')

    proj, hdu = load_projection("55.fits")

    wcs_in = WCS(proj.header)
    wcs_out = wcs_in.deepcopy()
    wcs_out.wcs.ctype = ['GLON-SIN', 'GLAT-SIN']
    wcs_out.wcs.crval = [134.37608, -31.939241]
    wcs_out.wcs.crpix = [2., 2.]

    header_out = proj.header
    header_out['NAXIS1'] = 4
    header_out['NAXIS2'] = 5
    header_out.update(wcs_out.to_header())

    result = proj.reproject(header_out)

    assert result.shape == (5, 4)
    assert result.beam == proj.beam

def test_nocelestial_reproject_2D_fail():

    pytest.importorskip('reproject')

    cube, data = cube_and_raw('255_delta.fits')

    proj = cube.moment0(axis=1)

    with pytest.raises(WCSCelestialError) as exc:
        proj.reproject(cube.header)

    assert exc.value.args[0] == ("WCS does not contain two spatial axes.")

@pytest.mark.parametrize('use_memmap', (True,False))
def test_downsample(use_memmap):
    cube, data = cube_and_raw('255.fits')

    dscube = cube.downsample_axis(factor=2, axis=0, use_memmap=use_memmap)

    expected = data.mean(axis=0)

    np.testing.assert_almost_equal(expected[None,:,:],
                                   dscube.filled_data[:].value)

    dscube = cube.downsample_axis(factor=2, axis=1, use_memmap=use_memmap)

    expected = np.array([data[:,:2,:].mean(axis=1),
                         data[:,2:4,:].mean(axis=1),
                         data[:,4:,:].mean(axis=1), # just data[:,4,:]
                        ]).swapaxes(0,1)
    assert expected.shape == (2,3,5)
    assert dscube.shape == (2,3,5)

    np.testing.assert_almost_equal(expected,
                                   dscube.filled_data[:].value)

    dscube = cube.downsample_axis(factor=2, axis=1, truncate=True,
                                  use_memmap=use_memmap)

    expected = np.array([data[:,:2,:].mean(axis=1),
                         data[:,2:4,:].mean(axis=1),
                        ]).swapaxes(0,1)

    np.testing.assert_almost_equal(expected,
                                   dscube.filled_data[:].value)
<commit_msg>add test for memmap reprojection<commit_after>import pytest
import numpy as np

from astropy import units as u
from astropy import convolution
from astropy.wcs import WCS
from astropy import wcs
from astropy.io import fits

from radio_beam import beam, Beam

from .. import SpectralCube
from ..utils import WCSCelestialError
from .test_spectral_cube import cube_and_raw
from .test_projection import load_projection
from . import path

def test_convolution():
    cube, data = cube_and_raw('255_delta.fits')

    # 1" convolved with 1.5" -> 1.8027....
    target_beam = Beam(1.802775637731995*u.arcsec, 1.802775637731995*u.arcsec,
                       0*u.deg)

    conv_cube = cube.convolve_to(target_beam)

    expected = convolution.Gaussian2DKernel((1.5*u.arcsec /
                                             beam.SIGMA_TO_FWHM /
                                             (5.555555555555e-4*u.deg)).decompose().value,
                                            x_size=5, y_size=5,
                                           )

    expected.normalize()

    np.testing.assert_almost_equal(expected.array,
                                   conv_cube.filled_data[0,:,:].value)

    # 2nd layer is all zeros
    assert np.all(conv_cube.filled_data[1,:,:] == 0.0)

def test_beams_convolution():
    cube, data = cube_and_raw('455_delta_beams.fits')

    # 1" convolved with 1.5" -> 1.8027....
    target_beam = Beam(1.802775637731995*u.arcsec, 1.802775637731995*u.arcsec,
                       0*u.deg)

    conv_cube = cube.convolve_to(target_beam)

    pixscale = wcs.utils.proj_plane_pixel_area(cube.wcs.celestial)**0.5*u.deg

    for ii, bm in enumerate(cube.beams):
        expected = target_beam.deconvolve(bm).as_kernel(pixscale, x_size=5,
                                                        y_size=5)
        expected.normalize()

        np.testing.assert_almost_equal(expected.array,
                                       conv_cube.filled_data[ii,:,:].value)

def test_beams_convolution_equal():
    cube, data = cube_and_raw('522_delta_beams.fits')

    # Only checking that the equal beam case is handled correctly.
    # Fake the beam in the first channel. Then ensure that the first channel
    # has NOT been convolved.
    target_beam = Beam(1.0 * u.arcsec, 1.0 * u.arcsec, 0.0 * u.deg)
    cube.beams.major[0] = target_beam.major
    cube.beams.minor[0] = target_beam.minor
    cube.beams.pa[0] = target_beam.pa

    conv_cube = cube.convolve_to(target_beam)

    np.testing.assert_almost_equal(cube.filled_data[0].value,
                                   conv_cube.filled_data[0].value)

@pytest.mark.parametrize('use_memmap', (True, False)
def test_reproject(use_memmap):

    pytest.importorskip('reproject')

    cube, data = cube_and_raw('adv.fits')

    wcs_in = WCS(cube.header)
    wcs_out = wcs_in.deepcopy()
    wcs_out.wcs.ctype = ['GLON-SIN', 'GLAT-SIN', wcs_in.wcs.ctype[2]]
    wcs_out.wcs.crval = [134.37608, -31.939241, wcs_in.wcs.crval[2]]
    wcs_out.wcs.crpix = [2., 2., wcs_in.wcs.crpix[2]]

    header_out = cube.header
    header_out['NAXIS1'] = 4
    header_out['NAXIS2'] = 5
    header_out['NAXIS3'] = cube.shape[0]
    header_out.update(wcs_out.to_header())

    result = cube.reproject(header_out, use_memmap=use_memmap)

    assert result.shape == (cube.shape[0], 5, 4)

def test_spectral_smooth():

    cube, data = cube_and_raw('522_delta.fits')

    result = cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0), use_memmap=False)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   convolution.Gaussian1DKernel(1.0,
                                                                x_size=5).array,
                                   4)

    result = cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0), use_memmap=True)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   convolution.Gaussian1DKernel(1.0,
                                                                x_size=5).array,
                                   4)

def test_spectral_smooth_4cores():

    pytest.importorskip('joblib')

    cube, data = cube_and_raw('522_delta.fits')

    result = cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0), num_cores=4, use_memmap=True)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   convolution.Gaussian1DKernel(1.0,
                                                                x_size=5).array,
                                   4)

    # this is one way to test non-parallel mode
    result = cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0), num_cores=4, use_memmap=False)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   convolution.Gaussian1DKernel(1.0,
                                                                x_size=5).array,
                                   4)

    # num_cores = 4 is a contradiction with parallel=False, but we want to make
    # sure it does the same thing
    result = cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0), num_cores=4, parallel=False)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   convolution.Gaussian1DKernel(1.0,
                                                                x_size=5).array,
                                   4)

def test_spectral_smooth_fail():

    cube, data = cube_and_raw('522_delta_beams.fits')

    with pytest.raises(AttributeError) as exc:
        cube.spectral_smooth(kernel=convolution.Gaussian1DKernel(1.0))

    assert exc.value.args[0] == ("VaryingResolutionSpectralCubes can't be "
                                 "spectrally smoothed.  Convolve to a "
                                 "common resolution with `convolve_to` before "
                                 "attempting spectral smoothed.")

def test_spectral_interpolate():

    cube, data = cube_and_raw('522_delta.fits')

    orig_wcs = cube.wcs.deepcopy()

    # midpoint between each position
    sg = (cube.spectral_axis[1:] + cube.spectral_axis[:-1])/2.

    result = cube.spectral_interpolate(spectral_grid=sg)

    np.testing.assert_almost_equal(result[:,0,0].value,
                                   [0.0, 0.5, 0.5, 0.0])

    assert cube.wcs.wcs.compare(orig_wcs.wcs)

def test_spectral_interpolate_with_fillvalue():

    cube, data = cube_and_raw('522_delta.fits')

    # Step one channel out of bounds.
    sg = ((cube.spectral_axis[0]) -
          (cube.spectral_axis[1] - cube.spectral_axis[0]) *
          np.linspace(1,4,4))
    result = cube.spectral_interpolate(spectral_grid=sg,
                                       fill_value=42)
    np.testing.assert_almost_equal(result[:,0,0].value,
                                   np.ones(4)*42)

def test_spectral_interpolate_fail():

    cube, data = cube_and_raw('522_delta_beams.fits')

    with pytest.raises(AttributeError) as exc:
        cube.spectral_interpolate(5)

    assert exc.value.args[0] == ("VaryingResolutionSpectralCubes can't be "
                                 "spectrally interpolated.  Convolve to a "
                                 "common resolution with `convolve_to` before "
                                 "attempting spectral interpolation.")

def test_spectral_interpolate_with_mask():

    hdu = fits.open(path("522_delta.fits"))[0]

    # Swap the velocity axis so indiff < 0 in spectral_interpolate
    hdu.header["CDELT3"] = - hdu.header["CDELT3"]

    cube = SpectralCube.read(hdu)

    mask = np.ones(cube.shape, dtype=bool)
    mask[:2] = False

    masked_cube = cube.with_mask(mask)

    orig_wcs = cube.wcs.deepcopy()

    # midpoint between each position
    sg = (cube.spectral_axis[1:] + cube.spectral_axis[:-1])/2.

    result = masked_cube.spectral_interpolate(spectral_grid=sg[::-1])

    # The output makes CDELT3 > 0 (reversed spectral axis) so the masked
    # portion are the final 2 channels.
    np.testing.assert_almost_equal(result[:, 0, 0].value,
                                   [0.0, 0.5, np.NaN, np.NaN])

    assert cube.wcs.wcs.compare(orig_wcs.wcs)

def test_spectral_interpolate_reversed():

    cube, data = cube_and_raw('522_delta.fits')

    orig_wcs = cube.wcs.deepcopy()

    # Reverse spectral axis
    sg = cube.spectral_axis[::-1]

    result = cube.spectral_interpolate(spectral_grid=sg)

    np.testing.assert_almost_equal(sg.value, result.spectral_axis.value)

def test_convolution_2D():

    proj, hdu = load_projection("55_delta.fits")

    # 1" convolved with 1.5" -> 1.8027....
    target_beam = Beam(1.802775637731995*u.arcsec, 1.802775637731995*u.arcsec,
                       0*u.deg)

    conv_proj = proj.convolve_to(target_beam)

    expected = convolution.Gaussian2DKernel((1.5*u.arcsec /
                                             beam.SIGMA_TO_FWHM /
                                             (5.555555555555e-4*u.deg)).decompose().value,
                                            x_size=5, y_size=5,
                                           )
    expected.normalize()

    np.testing.assert_almost_equal(expected.array,
                                   conv_proj.value)
    assert conv_proj.beam == target_beam

def test_nocelestial_convolution_2D_fail():

    cube, data = cube_and_raw('255_delta.fits')

    proj = cube.moment0(axis=1)

    test_beam = Beam(1.0 * u.arcsec)

    with pytest.raises(WCSCelestialError) as exc:
        proj.convolve_to(test_beam)

    assert exc.value.args[0] == ("WCS does not contain two spatial axes.")

def test_reproject_2D():

    pytest.importorskip('reproject')

    proj, hdu = load_projection("55.fits")

    wcs_in = WCS(proj.header)
    wcs_out = wcs_in.deepcopy()
    wcs_out.wcs.ctype = ['GLON-SIN', 'GLAT-SIN']
    wcs_out.wcs.crval = [134.37608, -31.939241]
    wcs_out.wcs.crpix = [2., 2.]

    header_out = proj.header
    header_out['NAXIS1'] = 4
    header_out['NAXIS2'] = 5
    header_out.update(wcs_out.to_header())

    result = proj.reproject(header_out)

    assert result.shape == (5, 4)
    assert result.beam == proj.beam

def test_nocelestial_reproject_2D_fail():

    pytest.importorskip('reproject')

    cube, data = cube_and_raw('255_delta.fits')

    proj = cube.moment0(axis=1)

    with pytest.raises(WCSCelestialError) as exc:
        proj.reproject(cube.header)

    assert exc.value.args[0] == ("WCS does not contain two spatial axes.")

@pytest.mark.parametrize('use_memmap', (True,False))
def test_downsample(use_memmap):
    cube, data = cube_and_raw('255.fits')

    dscube = cube.downsample_axis(factor=2, axis=0, use_memmap=use_memmap)

    expected = data.mean(axis=0)

    np.testing.assert_almost_equal(expected[None,:,:],
                                   dscube.filled_data[:].value)

    dscube = cube.downsample_axis(factor=2, axis=1, use_memmap=use_memmap)

    expected = np.array([data[:,:2,:].mean(axis=1),
                         data[:,2:4,:].mean(axis=1),
                         data[:,4:,:].mean(axis=1), # just data[:,4,:]
                        ]).swapaxes(0,1)
    assert expected.shape == (2,3,5)
    assert dscube.shape == (2,3,5)

    np.testing.assert_almost_equal(expected,
                                   dscube.filled_data[:].value)

    dscube = cube.downsample_axis(factor=2, axis=1, truncate=True,
                                  use_memmap=use_memmap)

    expected = np.array([data[:,:2,:].mean(axis=1),
                         data[:,2:4,:].mean(axis=1),
                        ]).swapaxes(0,1)

    np.testing.assert_almost_equal(expected,
                                   dscube.filled_data[:].value)
<|endoftext|>


<commit_before>import json
import unittest

import mock

import server
from server.models import sqldb
from server.studyspaces.models import StudySpacesBooking

class StudySpacesApiTests(unittest.TestCase):
    def setUp(self):
        server.app.config['TESTING'] = True
        sqldb.create_all()

    def tearDown(self):
        sqldb.session.remove()
        sqldb.drop_all()

    def testStudyspaceBooking(self):
        with server.app.test_client() as c:
            # fake the actual booking
            with mock.patch('penn.studyspaces.StudySpaces.book_room',
                            return_value={'success': 'booking placed', 'results': True}):
                resp = c.post('/studyspaces/book', data={
                    'building': 1,
                    'room': 1,
                    'start': '2017-02-08 10:00:00',
                    'end': '2017-02-08 10:30:00',
                    'firstname': 'Test',
                    'lastname': 'Test',
                    'email': 'test@example.com',
                    'groupname': 'Testing',
                    'phone': '000-000-0000',
                    'size': 1
                }, headers={'X-Device-Id': 'test'})
            res = json.loads(resp.data.decode('utf8'))
            self.assertTrue(len(res) > 0)

            # make sure the booking is saved to the database
            self.assertEquals(sqldb.session.query(StudySpacesBooking).count(), 1)

    def testStudyspaceCancelFailure(self):
        """Booking cancellation should not succeed if it is not in our database."""

        with server.app.test_client() as c:
            resp = c.post('/studyspaces/cancel', data={
                'booking_id': 'definitely_not_a_valid_booking_id_123'
            })
            res = json.loads(resp.data.decode('utf8'))
            self.assertTrue('error' in res)
<commit_msg>fixed flake<commit_after>import json
import unittest

import mock

import server
from server.models import sqldb
from server.studyspaces.models import StudySpacesBooking

class StudySpacesApiTests(unittest.TestCase):
    def setUp(self):
        server.app.config['TESTING'] = True
        sqldb.create_all()

    def tearDown(self):
        sqldb.session.remove()
        sqldb.drop_all()

    def testStudyspaceBooking(self):
        with server.app.test_client() as c:
            # fake the actual booking
            with mock.patch('penn.studyspaces.StudySpaces.book_room',
                            return_value={'success': 'booking placed', 'results': True}):
                resp = c.post('/studyspaces/book', data={
                    'building': 1,
                    'room': 1,
                    'start': '2017-02-08 10:00:00',
                    'end': '2017-02-08 10:30:00',
                    'firstname': 'Test',
                    'lastname': 'Test',
                    'email': 'test@example.com',
                    'groupname': 'Testing',
                    'phone': '000-000-0000',
                    'size': 1
                }, headers={'X-Device-Id': 'test'})
            res = json.loads(resp.data.decode('utf8'))
            self.assertTrue(len(res) > 0)

            # make sure the booking is saved to the database
            self.assertEquals(sqldb.session.query(StudySpacesBooking).count(), 1)

    def testStudyspaceCancelFailure(self):
        """Booking cancellation should not succeed if it is not in our database."""

        with server.app.test_client() as c:
            resp = c.post('/studyspaces/cancel', data={
                'booking_id': 'definitely_not_a_valid_booking_id_123'
            })
            res = json.loads(resp.data.decode('utf8'))
            self.assertTrue('error' in res)
<|endoftext|>


<commit_before># Copyright 2016 Goekcen Eraslan
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import pickle, os

import numpy as np
import pandas as pd
from sklearn.model_selection import KFold, train_test_split
from keras.models import load_model as keras_load_model

def read_csv(inputfile):
    matrix = pd.read_csv(inputfile, sep=None, engine='python')
    matrix = matrix.as_matrix()
    return matrix

def read_from_file(inputfile):
        return pickle.load(open(inputfile, 'rb'))

def load_model(log_dir):
    return keras_load_model("%s/weights.hdf5" % log_dir)

def preprocess(inputfile, kfold, transpose=False, outputfile=None,
               censorfile=None, censortranspose=False):
    matrix = read_csv(inputfile)
    if transpose:
        matrix = matrix.transpose()

    X = dict()

    if kfold:
        X['type'] = 'kfold'
        X['k'] = kfold
        X['folds'] = list()

        nsample = matrix.shape[0]

        # Prepare indices for k-fold cv and train/valid/test split
        # For example 5-fold generates 3-1-1 equally sized folds for
        # train/val/test tests
        for cv_trainval, cv_test in KFold(kfold, True, 42).split(range(nsample)):
            cv_train, cv_val = train_test_split(cv_trainval,
                    test_size=1/(kfold-1), random_state=42)
            X['folds'].append((matrix[cv_train, :],
                               matrix[cv_val,   :],
                               matrix[cv_test,  :]))

    else:
        X['type'] = 'valtest'
        X['train'], X['test'] = train_test_split(matrix, test_size=.1,
                                                 random_state=42)
        X['train'], X['val']  = train_test_split(X['train'], test_size=.1,
                                                 random_state=43)

    if censorfile:
        X['censor'] = read_csv(censorfile)
        if censortranspose:
            X['censor'] = X['censor'].transpose()

    if outputfile:
        # add file extension if missing
        outfile, outfile_ext = os.path.splitext(outputfile)
        if not outfile_ext:
            ext = 'ae'
            outputfile = '.'.join([outfile, ext])

        with open(outputfile, 'wb') as out:
            pickle.dump(X, out)

    return X

def preprocess_with_args(args):
    preprocess(args.input, kfold=args.kfold, transpose=args.transpose,
               outputfile=args.output, censorfile=args.censorfile,
               censortranspose=args.censortranspose)
<commit_msg>Compare input and censor file shapes<commit_after># Copyright 2016 Goekcen Eraslan
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import pickle, os

import numpy as np
import pandas as pd
from sklearn.model_selection import KFold, train_test_split
from keras.models import load_model as keras_load_model

def read_csv(inputfile):
    matrix = pd.read_csv(inputfile, sep=None, engine='python')
    matrix = matrix.as_matrix()
    return matrix

def read_from_file(inputfile):
        return pickle.load(open(inputfile, 'rb'))

def load_model(log_dir):
    return keras_load_model("%s/weights.hdf5" % log_dir)

def preprocess(inputfile, kfold, transpose=False, outputfile=None,
               censorfile=None, censortranspose=False):
    matrix = read_csv(inputfile)
    if transpose:
        matrix = matrix.transpose()

    X = dict()
    X['shape'] = matrix.shape

    if kfold:
        X['type'] = 'kfold'
        X['k'] = kfold
        X['folds'] = list()

        nsample = matrix.shape[0]

        # Prepare indices for k-fold cv and train/valid/test split
        # For example 5-fold generates 3-1-1 equally sized folds for
        # train/val/test tests
        for cv_trainval, cv_test in KFold(kfold, True, 42).split(range(nsample)):
            cv_train, cv_val = train_test_split(cv_trainval,
                    test_size=1/(kfold-1), random_state=42)
            X['folds'].append((matrix[cv_train, :],
                               matrix[cv_val,   :],
                               matrix[cv_test,  :]))

    else:
        X['type'] = 'valtest'
        X['train'], X['test'] = train_test_split(matrix, test_size=.1,
                                                 random_state=42)
        X['train'], X['val']  = train_test_split(X['train'], test_size=.1,
                                                 random_state=43)

    if censorfile:
        X['censor'] = read_csv(censorfile)
        if censortranspose:
            X['censor'] = X['censor'].transpose()

    assert X['censor'].shape == X['shape'], 'Input size of censorfile does not ' \
                                            'match the input file'

    if outputfile:
        # add file extension if missing
        outfile, outfile_ext = os.path.splitext(outputfile)
        if not outfile_ext:
            ext = 'ae'
            outputfile = '.'.join([outfile, ext])

        with open(outputfile, 'wb') as out:
            pickle.dump(X, out)

    return X

def preprocess_with_args(args):
    preprocess(args.input, kfold=args.kfold, transpose=args.transpose,
               outputfile=args.output, censorfile=args.censorfile,
               censortranspose=args.censortranspose)
<|endoftext|>


<commit_before>import pycept
import unittest

class ClientTestCase(unittest.TestCase):

  def testTokenizeEmptyString(self):
    client = pycept.Cept("foo", "bar")
    result = client.tokenize("")
    self.assertListEqual(result, [])

  def testTokenizeSingleSentence(self):
    client = pycept.Cept("foo", "bar")
    result = client.tokenize("The cow jumped over the moon.")
    self.assertEqual(len(result), 1)
    self.assertListEqual(result[0], ["cow", "jumped", "moon"])

  def testTokenizeMultipleSentences(self):
    client = pycept.Cept("foo", "bar")
    result = client.tokenize(
        "The cow jumped over the moon. And then the sun came up.")
    self.assertEqual(len(result), 2)
    self.assertListEqual(result[0], ["cow", "jumped", "moon"])
    self.assertListEqual(result[1], ["sun", "came"])

  def testTinyEmptyBitMapToSdr(self):
    client = pycept.Cept('foo', 'bar')
    result = client._bitmapToSdr({
      'width': 1,
      'height': 1,
      'positions': []
    })
    self.assertEqual(result, "0")

  def testTinyFullBitMapToSdr(self):
    client = pycept.Cept('foo', 'bar')
    result = client._bitmapToSdr({
      'width': 1,
      'height': 1,
      'positions': [0]
    })
    self.assertEqual(result, "1")

  def testSmallEmptyBitMapToSdr(self):
    client = pycept.Cept('foo', 'bar')
    result = client._bitmapToSdr({
      'width': 4,
      'height': 4,
      'positions': []
    })
    self.assertEqual(result, "0000000000000000")

  def testSmallFullBitMapToSdr(self):
    client = pycept.Cept('foo', 'bar')
    result = client._bitmapToSdr({
      'width': 4,
      'height': 4,
      'positions': [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
    })
    self.assertEqual(result, "1111111111111111")

  def testSmallPartialBitMapToSdr(self):
    client = pycept.Cept('foo', 'bar')
    result = client._bitmapToSdr({
      'width': 4,
      'height': 4,
      'positions': [2,3,6,7,10,11,14,15]
    })
    self.assertEqual(result, "0011001100110011")

if __name__ == "__main__":
  unittest.main()
<commit_msg>change file mode for nosetest discovery<commit_after>import pycept
import unittest

class ClientTestCase(unittest.TestCase):

  def testTokenizeEmptyString(self):
    client = pycept.Cept("foo", "bar")
    result = client.tokenize("")
    self.assertListEqual(result, [])

  def testTokenizeSingleSentence(self):
    client = pycept.Cept("foo", "bar")
    result = client.tokenize("The cow jumped over the moon.")
    self.assertEqual(len(result), 1)
    self.assertListEqual(result[0], ["cow", "jumped", "moon"])

  def testTokenizeMultipleSentences(self):
    client = pycept.Cept("foo", "bar")
    result = client.tokenize(
        "The cow jumped over the moon. And then the sun came up.")
    self.assertEqual(len(result), 2)
    self.assertListEqual(result[0], ["cow", "jumped", "moon"])
    self.assertListEqual(result[1], ["sun", "came"])

  def testTinyEmptyBitMapToSdr(self):
    client = pycept.Cept('foo', 'bar')
    result = client._bitmapToSdr({
      'width': 1,
      'height': 1,
      'positions': []
    })
    self.assertEqual(result, "0")

  def testTinyFullBitMapToSdr(self):
    client = pycept.Cept('foo', 'bar')
    result = client._bitmapToSdr({
      'width': 1,
      'height': 1,
      'positions': [0]
    })
    self.assertEqual(result, "1")

  def testSmallEmptyBitMapToSdr(self):
    client = pycept.Cept('foo', 'bar')
    result = client._bitmapToSdr({
      'width': 4,
      'height': 4,
      'positions': []
    })
    self.assertEqual(result, "0000000000000000")

  def testSmallFullBitMapToSdr(self):
    client = pycept.Cept('foo', 'bar')
    result = client._bitmapToSdr({
      'width': 4,
      'height': 4,
      'positions': [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
    })
    self.assertEqual(result, "1111111111111111")

  def testSmallPartialBitMapToSdr(self):
    client = pycept.Cept('foo', 'bar')
    result = client._bitmapToSdr({
      'width': 4,
      'height': 4,
      'positions': [2,3,6,7,10,11,14,15]
    })
    self.assertEqual(result, "0011001100110011")

if __name__ == "__main__":
  unittest.main()
<|endoftext|>


<commit_before>""" Copyright 2012, 2013 UW Information Technology, University of Washington

    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.

"""

from spotseeker_server.views.rest_dispatch import \
    RESTDispatch, RESTException, RESTFormInvalidError, JSONResponse
from spotseeker_server.forms.spot import SpotForm, SpotExtendedInfoForm
from spotseeker_server.default_forms.item import DefaultItemForm as ItemForm
from spotseeker_server.default_forms.item import DefaultItemExtendedInfoForm \
                                                as ItemExtendedInfoForm
from spotseeker_server.models import *
from django.http import HttpResponse
from spotseeker_server.require_auth import *
from django.db import transaction
import simplejson as json
import django.dispatch
from spotseeker_server.dispatch import \
    spot_pre_build, spot_pre_save, spot_post_save, spot_post_build
from django.core.exceptions import ObjectDoesNotExist

class ItemStash(object):
    """
    This object handles the storing and validation of an Item
    """
    def __init__(self, item):

        self.json = item
        item['item_category'] = item.pop('category', None)
        item['item_subcategory'] = item.pop('subcategory', None)
        self.form = ItemForm(item)

        if not self.form.is_valid():
            raise RESTFormInvalidError(self.form)

        if 'extended_info' not in item:
            raise ValidationError('extended_info required for item!')

        ei_json = self.json['extended_info']

        self.ei_forms = []

        for key, value in ei_json.iteritems():
            ei = {
                "key": key,
                "value": value
            }

            ei_form = ItemExtendedInfoForm(ei)
            self.ei_forms.append(ei_form)

            if not ei_form.is_valid():
                raise RESTFormInvalidError(ei_form)

    def get_form(self):
        """Returns the ItemForm used to validate the item"""
        return self.form

    def get_ei_forms(self):
        """Returns the list containing this Item's EI forms"""
        return self.ei_forms

    def remove_ei_form(self, form):
        """Removes a given EI form from ei_forms."""
        if form in self.ei_forms:
            self.ei_forms.remove(form)

    def get_json(self):
        return self.json

    def set_instance(self, instance):
        self.instance = instance
        self.form = ItemForm(self.json, instance=instance)

    def get_instance(self):
        return self.instance

@django.dispatch.receiver(
    spot_pre_build,
    dispatch_uid='spotseeker_server.views.spot.stash_items')
def _stash_items(sender, **kwargs):
    """Save the available hours for later"""
    json_values = kwargs['json_values']
    stash = kwargs['stash']

    if 'items' not in json_values:
        return

    json_items = json_values['items']

    stash['updated_items'] = []
    stash['new_items'] = []

    # create the items
    for item in json_items:
        # put the item in either new or updated items
        if 'id' in item:
            stash['updated_items'].append(ItemStash(item))
        else:
            stash['new_items'].append(ItemStash(item))

@django.dispatch.receiver(
    spot_pre_save,
    dispatch_uid='spotseeker_server.views.spot.clean_updated_items')
def _clean_updated_items(sender, **kwargs):
    """Removes unchanged items/item EI and sets items to delete"""
    json_values = kwargs['json_values']
    stash = kwargs['stash']
    spot = kwargs['spot']

    stash['items_to_delete'] = []
    stash['item_ei_to_delete'] = []

    # if we don't have any json items, then return
    if spot is None or 'updated_items' not in stash:
        return

    updated_items = stash['updated_items']

    updated_item_models = []
    updated_item_models_to_stash = {}

    # create the lists of items to delete
    items_to_delete = []
    item_ei_to_delete = []

    # get the old items
    old_items = spot.item_set.all()

    # create item models so we can use a hashmap for matching
    for item in updated_items:
        item_json = item.get_json()
        item_model = Item(name=item_json['name'],
                          item_category=item_json['item_category'],
                          item_subcategory=item_json['item_subcategory'],
                          id=item_json['id'],
                          spot=spot)
        updated_item_models.append(item_model)
        updated_item_models_to_stash[item_model] = item

    # create a hashmap to match old to new, by using old:old
    lookup_hashmap = {}

    for old_item in old_items:
        lookup_hashmap[old_item] = old_item

    equality_hashmap = {}
    # create a hashmap matching new to old
    for updated_item in updated_item_models:
        if updated_item in lookup_hashmap:
            equality_hashmap[updated_item] = lookup_hashmap.pop(updated_item)

    # we should delete any missing from the PUT
    for item_to_delete in lookup_hashmap:
        items_to_delete.append(item_to_delete)

    # find items that haven't been updated and remove them
    for updated_item_model, old_item, in equality_hashmap.iteritems():

        updated_item = updated_item_models_to_stash[updated_item_model]
        updated_item_form = updated_item.get_form()
        updated_item.set_instance(old_item)
        updated_item_ei = updated_item.get_ei_forms()

        # clean up the EI
        old_ei_set = old_item.itemextendedinfo_set.all()
        ei_to_remove = clean_ei(old_ei_set, updated_item_ei)
        item_ei_to_delete += ei_to_remove

        # get rid of items that are all the same without EI
        if (updated_item_model.name == old_item.name and
            updated_item_model.item_category == old_item.item_category and
            updated_item_model.item_subcategory ==
            old_item.item_subcategory and
                len(updated_item_ei) == 0):
            updated_items.remove(updated_item)

    stash['items_to_delete'] = items_to_delete
    stash['item_ei_to_delete'] = item_ei_to_delete

def clean_ei(old_ei_list, new_ei_forms):
    """Returns deleted EI keys and removes unchanged EI"""
    to_delete = []
    forms_to_remove = []
    for old_ei in old_ei_list:
        found = True

        for ei_form in new_ei_forms:
            if ei_form.cleaned_data['key'] != old_ei.key:
                found = False

        # put the EI that has been removed in to_remove for removal
        if found:
            to_delete.append(old_ei)

    # get rid of unchanged item EI
    for form in forms_to_remove:
        new_ei_forms.remove(form)

    return to_delete

@django.dispatch.receiver(
    spot_post_save,
    dispatch_uid='spotseeker_server.views.spot.save_items')
def _save_items(sender, **kwargs):
    """Save updated and new items, and then delete items not included in PUT"""
    spot = kwargs['spot']
    stash = kwargs['stash']

    if ('new_items' not in stash or 'updated_items' not in stash or
            'items_to_delete' not in stash or
            'item_ei_to_delete' not in stash):
        return

    new_items = stash['new_items']
    updated_items = stash['updated_items']

    # combine new_items and updated_items now that the middleware has handled
    # them
    for item in updated_items:
        new_items.append(item)

    # save the new item, and set the spot
    for item in new_items:
        item_form = item.get_form()
        item_model = item_form.save(commit=False)
        item_model.spot = spot

        item_model.save()

        ei_forms = item.get_ei_forms()

        for item_ei in ei_forms:
            # save the new EI
            item_ei_model = item_ei.save(commit=False)
            try:
                actual_item_ei = item_model.itemextendedinfo_set.get(key=item_ei_model.key)
                if actual_item_ei.value != item_ei_model.value:
                    actual_item_ei.value = item_ei_model.value
                    actual_item_ei.save()
            except ObjectDoesNotExist:
                item_ei_model.item = item_model
                item_ei_model.save()

    # delete items not included
    items_to_delete = stash['items_to_delete']

    for item in items_to_delete:
        item.delete()

    if 'item_ei_to_delete' in stash:
        item_ei_to_delete = stash['item_ei_to_delete']

        for ei in item_ei_to_delete:
            ei.delete()
<commit_msg>Revise how removed ItemExtendedInfo is deleted. Fixes failing unit test spotseeker_server.test.spot_put.SpotPUTTest.test_remove_item_extended_info<commit_after>""" Copyright 2012, 2013 UW Information Technology, University of Washington

    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.

"""

from spotseeker_server.views.rest_dispatch import \
    RESTDispatch, RESTException, RESTFormInvalidError, JSONResponse
from spotseeker_server.forms.spot import SpotForm, SpotExtendedInfoForm
from spotseeker_server.default_forms.item import DefaultItemForm as ItemForm
from spotseeker_server.default_forms.item import DefaultItemExtendedInfoForm \
                                                as ItemExtendedInfoForm
from spotseeker_server.models import *
from django.http import HttpResponse
from spotseeker_server.require_auth import *
from django.db import transaction
import simplejson as json
import django.dispatch
from spotseeker_server.dispatch import \
    spot_pre_build, spot_pre_save, spot_post_save, spot_post_build
from django.core.exceptions import ObjectDoesNotExist

class ItemStash(object):
    """
    This object handles the storing and validation of an Item
    """
    def __init__(self, item):

        self.json = item
        item['item_category'] = item.pop('category', None)
        item['item_subcategory'] = item.pop('subcategory', None)
        self.form = ItemForm(item)

        if not self.form.is_valid():
            raise RESTFormInvalidError(self.form)

        if 'extended_info' not in item:
            raise ValidationError('extended_info required for item!')

        ei_json = self.json['extended_info']

        self.ei_forms = []

        for key, value in ei_json.iteritems():
            ei = {
                "key": key,
                "value": value
            }

            ei_form = ItemExtendedInfoForm(ei)
            self.ei_forms.append(ei_form)

            if not ei_form.is_valid():
                raise RESTFormInvalidError(ei_form)

    def get_form(self):
        """Returns the ItemForm used to validate the item"""
        return self.form

    def get_ei_forms(self):
        """Returns the list containing this Item's EI forms"""
        return self.ei_forms

    def remove_ei_form(self, form):
        """Removes a given EI form from ei_forms."""
        if form in self.ei_forms:
            self.ei_forms.remove(form)

    def get_json(self):
        return self.json

    def set_instance(self, instance):
        self.instance = instance
        self.form = ItemForm(self.json, instance=instance)

    def get_instance(self):
        return self.instance

@django.dispatch.receiver(
    spot_pre_build,
    dispatch_uid='spotseeker_server.views.spot.stash_items')
def _stash_items(sender, **kwargs):
    """Save the available hours for later"""
    json_values = kwargs['json_values']
    stash = kwargs['stash']

    if 'items' not in json_values:
        return

    json_items = json_values['items']

    stash['updated_items'] = []
    stash['new_items'] = []

    # create the items
    for item in json_items:
        # put the item in either new or updated items
        if 'id' in item:
            stash['updated_items'].append(ItemStash(item))
        else:
            stash['new_items'].append(ItemStash(item))

@django.dispatch.receiver(
    spot_pre_save,
    dispatch_uid='spotseeker_server.views.spot.clean_updated_items')
def _clean_updated_items(sender, **kwargs):
    """Removes unchanged items/item EI and sets items to delete"""
    json_values = kwargs['json_values']
    stash = kwargs['stash']
    spot = kwargs['spot']

    stash['items_to_delete'] = []
    stash['item_ei_to_delete'] = []

    # if we don't have any json items, then return
    if spot is None or 'updated_items' not in stash:
        return

    updated_items = stash['updated_items']

    updated_item_models = []
    updated_item_models_to_stash = {}

    # create the lists of items to delete
    items_to_delete = []
    item_ei_to_delete = []

    # get the old items
    old_items = spot.item_set.all()

    # create item models so we can use a hashmap for matching
    for item in updated_items:
        item_json = item.get_json()
        item_model = Item(name=item_json['name'],
                          item_category=item_json['item_category'],
                          item_subcategory=item_json['item_subcategory'],
                          id=item_json['id'],
                          spot=spot)
        updated_item_models.append(item_model)
        updated_item_models_to_stash[item_model] = item

    # create a hashmap to match old to new, by using old:old
    lookup_hashmap = {}

    for old_item in old_items:
        lookup_hashmap[old_item] = old_item

    equality_hashmap = {}
    # create a hashmap matching new to old
    for updated_item in updated_item_models:
        if updated_item in lookup_hashmap:
            equality_hashmap[updated_item] = lookup_hashmap.pop(updated_item)

    # we should delete any missing from the PUT
    for item_to_delete in lookup_hashmap:
        items_to_delete.append(item_to_delete)

    # find items that haven't been updated and remove them
    for updated_item_model, old_item, in equality_hashmap.iteritems():

        updated_item = updated_item_models_to_stash[updated_item_model]
        updated_item_form = updated_item.get_form()
        updated_item.set_instance(old_item)
        updated_item_ei = updated_item.get_ei_forms()

        # clean up the EI
        old_ei_set = old_item.itemextendedinfo_set.all()
        ei_to_remove = clean_ei(old_ei_set, updated_item_ei)
        item_ei_to_delete += ei_to_remove

        # get rid of items that are all the same without EI
        if (updated_item_model.name == old_item.name and
            updated_item_model.item_category == old_item.item_category and
            updated_item_model.item_subcategory ==
            old_item.item_subcategory and
                len(updated_item_ei) == 0):
            updated_items.remove(updated_item)

    stash['items_to_delete'] = items_to_delete
    stash['item_ei_to_delete'] = item_ei_to_delete

def clean_ei(old_ei_list, new_ei_forms):
    """Returns deleted EI keys and removes unchanged EI"""
    to_delete = []
    forms_to_remove = []
    new_ei_form_keys = []

    for ei_form in new_ei_forms:
        new_ei_form_keys.append(ei_form.cleaned_data['key'])

    for old_ei in old_ei_list:
        if old_ei.key not in new_ei_form_keys:
            to_delete.append(old_ei)

    return to_delete

@django.dispatch.receiver(
    spot_post_save,
    dispatch_uid='spotseeker_server.views.spot.save_items')
def _save_items(sender, **kwargs):
    """Save updated and new items, and then delete items not included in PUT"""
    spot = kwargs['spot']
    stash = kwargs['stash']

    if ('new_items' not in stash or 'updated_items' not in stash or
            'items_to_delete' not in stash or
            'item_ei_to_delete' not in stash):
        return

    new_items = stash['new_items']
    updated_items = stash['updated_items']

    # combine new_items and updated_items now that the middleware has handled
    # them
    for item in updated_items:
        new_items.append(item)

    # save the new item, and set the spot
    for item in new_items:
        item_form = item.get_form()
        item_model = item_form.save(commit=False)
        item_model.spot = spot

        item_model.save()

        ei_forms = item.get_ei_forms()

        for item_ei in ei_forms:
            # save the new EI
            item_ei_model = item_ei.save(commit=False)
            try:
                actual_item_ei = item_model.itemextendedinfo_set.get(key=item_ei_model.key)
                if actual_item_ei.value != item_ei_model.value:
                    actual_item_ei.value = item_ei_model.value
                    actual_item_ei.save()
            except ObjectDoesNotExist:
                item_ei_model.item = item_model
                item_ei_model.save()

    # delete items not included
    items_to_delete = stash['items_to_delete']

    for item in items_to_delete:
        item.delete()

    if 'item_ei_to_delete' in stash:
        item_ei_to_delete = stash['item_ei_to_delete']

        for ei in item_ei_to_delete:
            ei.delete()
<|endoftext|>