<commit_before>#!/usr/bin/env python

## TODO: DISCLAIMER, LICENSE, STUFF,...

import rospy

from frame_editor.objects import *
from frame_editor.commands import *
from frame_editor.interface import Interface

from frame_editor.constructors_geometry import *
from frame_editor.constructors_std import *

from geometry_msgs.msg import Pose

from interactive_markers.interactive_marker_server import *
from visualization_msgs.msg import InteractiveMarkerControl, Marker

class FrameEditor_InteractiveMarker(Interface):

    def __init__(self, frame_editor):
        self.editor = frame_editor
        self.editor.observers.append(self)

        self.server = InteractiveMarkerServer("frame_editor_interactive")

        self.set_marker_settings(["x", "y", "z", "a", "b", "c"])

        self.old_frame = None

    def update(self, editor, level, elements):

        if level & 2:
            ## Check for change
            if self.editor.active_frame is not self.old_frame:
                self.make_interactive(self.editor.active_frame)

        if level & 4:
            if self.editor.active_frame is not None:
                self.int_marker.name = self.editor.active_frame.name
                self.int_marker.header.frame_id = self.editor.active_frame.parent
                self.int_marker.pose = self.editor.active_frame.pose
                self.server.insert(self.int_marker, self.callback_marker)
                self.server.applyChanges()

    def make_interactive(self, frame):

        ## Stop currently active frame
        if self.old_frame is not None:
            self.server.erase(self.old_frame.name)
            self.server.applyChanges()

        self.old_frame = frame

        if frame is not None:
            self.set_marker_settings(["x", "y", "z", "a", "b", "c"], frame)

            self.int_marker.name = frame.name
            self.int_marker.header.frame_id = frame.parent
            self.int_marker.pose = frame.pose

            self.server.insert(self.int_marker, self.callback_marker)
            self.server.applyChanges()

    def callback_marker(self, feedback):
        self.editor.command(Command_SetPose(self.editor, self.editor.active_frame, FromPoint(feedback.pose.position), FromQuaternion(feedback.pose.orientation)))

    def set_marker_settings(self, arrows, frame=None, scale=0.25):
        '''arrows is a list with any number of the following strings: ["x", "y", "z", "a", "b", "c"].'''

        if frame:
            style = frame.style
        else:
            style = "none"

        ## Marker
        int_marker = InteractiveMarker()
        int_marker.header.frame_id = ""
        int_marker.name = ""
        int_marker.description = "Frame Editor"
        int_marker.pose = Pose()
        int_marker.scale = scale

        if "x" in arrows:
            control = InteractiveMarkerControl()
            control.name = "move_x"
            control.orientation = NewQuaternion(1, 0, 0, 1)
            control.interaction_mode = InteractiveMarkerControl.MOVE_AXIS
            int_marker.controls.append(control);

        if "y" in arrows:
            control = InteractiveMarkerControl()
            control.name = "move_y"
            control.orientation = NewQuaternion(0, 1, 0, 1)
            control.interaction_mode = InteractiveMarkerControl.MOVE_AXIS
            int_marker.controls.append(control);

        if "z" in arrows:
            control = InteractiveMarkerControl()
            control.name = "move_z"
            control.orientation = NewQuaternion(0, 0, 1, 1)
            control.interaction_mode = InteractiveMarkerControl.MOVE_AXIS
            int_marker.controls.append(control);

        if "a" in arrows:
            control = InteractiveMarkerControl()
            control.name = "rotate_x"
            control.orientation = NewQuaternion(1, 0, 0, 1)
            control.interaction_mode = InteractiveMarkerControl.ROTATE_AXIS
            int_marker.controls.append(control);

        if "b" in arrows:
            control = InteractiveMarkerControl()
            control.name = "rotate_y"
            control.orientation = NewQuaternion(0, 1, 0, 1)
            control.interaction_mode = InteractiveMarkerControl.ROTATE_AXIS
            int_marker.controls.append(control);

        if "c" in arrows:
            control = InteractiveMarkerControl()
            control.name = "rotate_z"
            control.orientation = NewQuaternion(0, 0, 1, 1)
            control.interaction_mode = InteractiveMarkerControl.ROTATE_AXIS
            int_marker.controls.append(control);

        ## Style ##
        ##
        style_marker = Marker()
        style_marker.scale = NewVector3(0.75*scale, 0.75*scale, 0.75*scale)
        style_marker.color = NewColor(0.0, 0.5, 0.5, 0.75)

        if style != "none":
            style_marker = frame.marker

        style_control = InteractiveMarkerControl()
        style_control.always_visible = True
        style_control.markers.append(style_marker)

        if style != "none":
            int_marker.controls.append(style_control)

        self.int_marker = int_marker

# eof
<commit_msg>frame editor: interactive marker doesn't show style marker any more (redundant). Use the intendet marker topic to show style markers<commit_after>#!/usr/bin/env python

## TODO: DISCLAIMER, LICENSE, STUFF,...

import rospy

from frame_editor.objects import *
from frame_editor.commands import *
from frame_editor.interface import Interface

from frame_editor.constructors_geometry import *
from frame_editor.constructors_std import *

from geometry_msgs.msg import Pose

from interactive_markers.interactive_marker_server import *
from visualization_msgs.msg import InteractiveMarkerControl, Marker

class FrameEditor_InteractiveMarker(Interface):

    def __init__(self, frame_editor):
        self.editor = frame_editor
        self.editor.observers.append(self)

        self.server = InteractiveMarkerServer("frame_editor_interactive")

        self.set_marker_settings(["x", "y", "z", "a", "b", "c"])

        self.old_frame = None

    def update(self, editor, level, elements):

        if level & 2:
            ## Check for change
            if self.editor.active_frame is not self.old_frame:
                self.make_interactive(self.editor.active_frame)

        if level & 4:
            if self.editor.active_frame is not None:
                self.int_marker.name = self.editor.active_frame.name
                self.int_marker.header.frame_id = self.editor.active_frame.parent
                self.int_marker.pose = self.editor.active_frame.pose
                self.server.insert(self.int_marker, self.callback_marker)
                self.server.applyChanges()

    def make_interactive(self, frame):

        ## Stop currently active frame
        if self.old_frame is not None:
            self.server.erase(self.old_frame.name)
            self.server.applyChanges()

        self.old_frame = frame

        if frame is not None:
            self.set_marker_settings(["x", "y", "z", "a", "b", "c"], frame)

            self.int_marker.name = frame.name
            self.int_marker.header.frame_id = frame.parent
            self.int_marker.pose = frame.pose

            self.server.insert(self.int_marker, self.callback_marker)
            self.server.applyChanges()

    def callback_marker(self, feedback):
        self.editor.command(Command_SetPose(self.editor, self.editor.active_frame, FromPoint(feedback.pose.position), FromQuaternion(feedback.pose.orientation)))

    def set_marker_settings(self, arrows, frame=None, scale=0.25):
        '''arrows is a list with any number of the following strings: ["x", "y", "z", "a", "b", "c"].'''

        if frame:
            style = frame.style
        else:
            style = "none"

        ## Marker
        int_marker = InteractiveMarker()
        int_marker.header.frame_id = ""
        int_marker.name = ""
        int_marker.description = "Frame Editor"
        int_marker.pose = Pose()
        int_marker.scale = scale

        if "x" in arrows:
            control = InteractiveMarkerControl()
            control.name = "move_x"
            control.orientation = NewQuaternion(1, 0, 0, 1)
            control.interaction_mode = InteractiveMarkerControl.MOVE_AXIS
            int_marker.controls.append(control);

        if "y" in arrows:
            control = InteractiveMarkerControl()
            control.name = "move_y"
            control.orientation = NewQuaternion(0, 1, 0, 1)
            control.interaction_mode = InteractiveMarkerControl.MOVE_AXIS
            int_marker.controls.append(control);

        if "z" in arrows:
            control = InteractiveMarkerControl()
            control.name = "move_z"
            control.orientation = NewQuaternion(0, 0, 1, 1)
            control.interaction_mode = InteractiveMarkerControl.MOVE_AXIS
            int_marker.controls.append(control);

        if "a" in arrows:
            control = InteractiveMarkerControl()
            control.name = "rotate_x"
            control.orientation = NewQuaternion(1, 0, 0, 1)
            control.interaction_mode = InteractiveMarkerControl.ROTATE_AXIS
            int_marker.controls.append(control);

        if "b" in arrows:
            control = InteractiveMarkerControl()
            control.name = "rotate_y"
            control.orientation = NewQuaternion(0, 1, 0, 1)
            control.interaction_mode = InteractiveMarkerControl.ROTATE_AXIS
            int_marker.controls.append(control);

        if "c" in arrows:
            control = InteractiveMarkerControl()
            control.name = "rotate_z"
            control.orientation = NewQuaternion(0, 0, 1, 1)
            control.interaction_mode = InteractiveMarkerControl.ROTATE_AXIS
            int_marker.controls.append(control);

        ## Style ##
        ##
        #style_marker = Marker()
        #style_marker.scale = NewVector3(0.75*scale, 0.75*scale, 0.75*scale)
        #style_marker.color = NewColor(0.0, 0.5, 0.5, 0.75)

        #if style != "none":
        #    style_marker = frame.marker

        #style_control = InteractiveMarkerControl()
        #style_control.always_visible = True
        #style_control.markers.append(style_marker)

        #if style != "none":
        #    int_marker.controls.append(style_control)

        self.int_marker = int_marker

# eof
<|endoftext|>


<commit_before># Ant-FS
#
# Copyright (c) 2012, Gustav Tiger <gustav@tiger.name>
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

import array
import logging
import struct
import threading
import Queue

from ant.easy.channel import Channel
from ant.easy.node import Node, Message

import ant.fs.command
from ant.fs.beacon import Beacon
from ant.fs.command import LinkCommand, DownloadRequest, DownloadResponse, \
        AuthenticateCommand, AuthenticateResponse, DisconnectCommand
from ant.fs.file import Directory, File

_logger = logging.getLogger("garmin.ant.fs.manager")

class AntFSException(Exception):
    pass

class AntFSDownloadException(AntFSException):
    pass

class AntFSAuthenticationException(AntFSException):

    def __init__(self, error):
        self.error = error

class Application:

    _serial_number = 1337
    _frequency     = 19    # 0 to 124, x - 2400 (in MHz)
    #_period        =

    def __init__(self):

        self._queue = Queue.Queue()
        self._beacons = Queue.Queue()

        self._node = Node(0x0fcf, 0x1008)

        print "Request basic information..."
        m = self._node.request_message(Message.ID.RESPONSE_VERSION)
        print "  ANT version:  ", struct.unpack("<10sx", m[2])[0]
        m = self._node.request_message(Message.ID.RESPONSE_CAPABILITIES)
        print "  Capabilities: ", m[2]
        m = self._node.request_message(Message.ID.RESPONSE_SERIAL_NUMBER)
        print "  Serial number:", struct.unpack("<I", m[2])[0]

        print "Starting system..."

        NETWORK_KEY= [0xa8, 0xa4, 0x23, 0xb9, 0xf5, 0x5e, 0x63, 0xc1]

        self._node.reset_system()
        self._node.set_network_key(0x00, NETWORK_KEY)

        self._channel = self._node.new_channel(Channel.Type.BIDIRECTIONAL_RECEIVE)
        self._channel.on_broadcast_data = self._on_data
        self._channel.on_burst_data = self._on_data

        self.setup_channel(self._channel)

        self._worker_thread =threading.Thread(target=self._worker, name="ant.fs")
        self._worker_thread.start()

    def _worker(self):
        self._node.start()

    def _main(self):
        #print "Run, await beacon"
        b = self._get_beacon()
        #print "Run, do link"
        self.on_link(b)
        #self.link()
        for i in range(0, 5):
            #print "Run, await beacon"
            x = self._get_beacon()
            #print "Run, check beacon"
            if x.get_client_device_state() == Beacon.ClientDeviceState.AUTHENTICATION:
                #print "Run, beacon ok"
                self.on_authentication(x)
                #self.authentication()
                y = self._get_beacon()
                self.on_transport(y)
                self.disconnect()
                break
            else:
                #print "Run, beacon bad :(", x.get_client_device_state(), Beacon.ClientDeviceState.AUTHENTICATION
                pass
        #else:
        #    run()

        # auth state
        _logger.debug("Run 5")
        self.stop()

    def _on_beacon(self, data):
        b = Beacon.parse(data)
        self._beacons.put(b)

    def _on_command(self, data):
        c = ant.fs.command.parse(data)
        self._queue.put(c)

    def _on_data(self, data):
        #print "_on_data", data, len(data)
        if data[0] == 0x43:
            self._on_beacon(data[:8])
            if len(data[8:]) > 0:
                self._on_command(data[8:])
        elif data[0] == 0x44:
            self._on_command(data)

    def _get_beacon(self):
        b = self._beacons.get()
        self._beacons.task_done()
        return b

    def _get_command(self, timeout=3.0):
        _logger.debug("Get command, t%d, s%d", timeout, self._queue.qsize())
        c = self._queue.get(True, timeout)
        self._queue.task_done()
        return c

    def _send_command(self, c):
        data = c.get()
        if len(data) == 8:
            self._channel.send_acknowledged_data(data)
        else:
            self._channel.send_burst_transfer(data)

    # Application actions are defined from here
    # =======================================================================

    # These should be overloaded:

    def setup_channel(self, channel):
        pass

    def on_link(self, beacon):
        pass

    def on_authentication(self, beacon):
        pass

    def on_transport(self, beacon):
        pass

    # Shouldn't have to touch these:

    def start(self):
        self._main()

    def stop(self):
        self._node.stop()

    def erase(self, index):
        pass

    def upload(self, index, data):
        pass

    def download(self, index, callback=None):
        offset  = 0
        initial = True
        crc     = 0
        data    = array.array('B')
        while True:
            _logger.debug("Download %d, o%d, c%d", index, offset, crc)
            self._send_command(DownloadRequest(index, offset, True, crc))
            _logger.debug("Wait for response...")
            try:
                response = self._get_command()
                if isinstance(response, DownloadResponse):
                    remaining    = response._get_argument("remaining")
                    offset       = response._get_argument("offset")
                    total        = offset + remaining
                    data[offset:total] = response._get_argument("data")[:remaining]
                    #print "rem", remaining, "offset", offset, "total", total, "size", response._get_argument("size")
                    if callback != None:
                        callback(float(total) / float(response._get_argument("size")))
                    if total == response._get_argument("size"):
                        return data
                    crc = response._get_argument("crc")
                    offset = total
                else:
                    raise AntFSDownloadException()
            except Queue.Empty:
                _logger.debug("Download %d timeout", index)
                #print "recover from download failure"

    def download_directory(self, callback=None):
        data = self.download(0, callback)
        return Directory.parse(data)

    def link(self):
        self._channel.request_message(Message.ID.RESPONSE_CHANNEL_ID)
        self._send_command(LinkCommand(self._frequency, 4, self._serial_number))

        # New period, search timeout
        self._channel.set_period(4096)
        self._channel.set_search_timeout(3)
        self._channel.set_rf_freq(self._frequency)

    def authentication_serial(self):
        self._send_command(AuthenticateCommand(
                AuthenticateCommand.Request.SERIAL,
                self._serial_number))
        response = self._get_command()
        return (response.get_serial(), response.get_data_string())

    def authentication_passkey(self, passkey):
        self._send_command(AuthenticateCommand(
                AuthenticateCommand.Request.PASSKEY_EXCHANGE,
                self._serial_number, passkey))

        response = self._get_command()
        if response._get_argument("type") == AuthenticateResponse.Response.ACCEPT:
            return response.get_data_array()
        else:
            raise AntFSAuthenticationException(response._get_argument("type"))

    def authentication_pair(self, friendly_name):
        data = array.array('B', map(ord, list(friendly_name)))
        self._send_command(AuthenticateCommand(
                AuthenticateCommand.Request.PAIRING,
                self._serial_number, data))

        response = self._get_command(30)
        if response._get_argument("type") == AuthenticateResponse.Response.ACCEPT:
            return response.get_data_array()
        else:
            raise AntFSAuthenticationException(response._get_argument("type"))

    def disconnect(self):
        d = DisconnectCommand(DisconnectCommand.Type.RETURN_LINK, 0, 0)
        self._send_command(d)

<commit_msg>Report download errors<commit_after># Ant-FS
#
# Copyright (c) 2012, Gustav Tiger <gustav@tiger.name>
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

import array
import logging
import struct
import threading
import Queue

from ant.easy.channel import Channel
from ant.easy.node import Node, Message

import ant.fs.command
from ant.fs.beacon import Beacon
from ant.fs.command import LinkCommand, DownloadRequest, DownloadResponse, \
        AuthenticateCommand, AuthenticateResponse, DisconnectCommand
from ant.fs.file import Directory, File

_logger = logging.getLogger("garmin.ant.fs.manager")

class AntFSException(Exception):
    pass

class AntFSDownloadException(AntFSException):

    def __init__(self, error):
        self.error = error

class AntFSAuthenticationException(AntFSException):

    def __init__(self, error):
        self.error = error

class Application:

    _serial_number = 1337
    _frequency     = 19    # 0 to 124, x - 2400 (in MHz)
    #_period        =

    def __init__(self):

        self._queue = Queue.Queue()
        self._beacons = Queue.Queue()

        self._node = Node(0x0fcf, 0x1008)

        print "Request basic information..."
        m = self._node.request_message(Message.ID.RESPONSE_VERSION)
        print "  ANT version:  ", struct.unpack("<10sx", m[2])[0]
        m = self._node.request_message(Message.ID.RESPONSE_CAPABILITIES)
        print "  Capabilities: ", m[2]
        m = self._node.request_message(Message.ID.RESPONSE_SERIAL_NUMBER)
        print "  Serial number:", struct.unpack("<I", m[2])[0]

        print "Starting system..."

        NETWORK_KEY= [0xa8, 0xa4, 0x23, 0xb9, 0xf5, 0x5e, 0x63, 0xc1]

        self._node.reset_system()
        self._node.set_network_key(0x00, NETWORK_KEY)

        self._channel = self._node.new_channel(Channel.Type.BIDIRECTIONAL_RECEIVE)
        self._channel.on_broadcast_data = self._on_data
        self._channel.on_burst_data = self._on_data

        self.setup_channel(self._channel)

        self._worker_thread =threading.Thread(target=self._worker, name="ant.fs")
        self._worker_thread.start()

    def _worker(self):
        self._node.start()

    def _main(self):
        #print "Run, await beacon"
        b = self._get_beacon()
        #print "Run, do link"
        self.on_link(b)
        #self.link()
        for i in range(0, 5):
            #print "Run, await beacon"
            x = self._get_beacon()
            #print "Run, check beacon"
            if x.get_client_device_state() == Beacon.ClientDeviceState.AUTHENTICATION:
                #print "Run, beacon ok"
                self.on_authentication(x)
                #self.authentication()
                y = self._get_beacon()
                self.on_transport(y)
                self.disconnect()
                break
            else:
                #print "Run, beacon bad :(", x.get_client_device_state(), Beacon.ClientDeviceState.AUTHENTICATION
                pass
        #else:
        #    run()

        # auth state
        _logger.debug("Run 5")
        self.stop()

    def _on_beacon(self, data):
        b = Beacon.parse(data)
        self._beacons.put(b)

    def _on_command(self, data):
        c = ant.fs.command.parse(data)
        self._queue.put(c)

    def _on_data(self, data):
        #print "_on_data", data, len(data)
        if data[0] == 0x43:
            self._on_beacon(data[:8])
            if len(data[8:]) > 0:
                self._on_command(data[8:])
        elif data[0] == 0x44:
            self._on_command(data)

    def _get_beacon(self):
        b = self._beacons.get()
        self._beacons.task_done()
        return b

    def _get_command(self, timeout=3.0):
        _logger.debug("Get command, t%d, s%d", timeout, self._queue.qsize())
        c = self._queue.get(True, timeout)
        self._queue.task_done()
        return c

    def _send_command(self, c):
        data = c.get()
        if len(data) == 8:
            self._channel.send_acknowledged_data(data)
        else:
            self._channel.send_burst_transfer(data)

    # Application actions are defined from here
    # =======================================================================

    # These should be overloaded:

    def setup_channel(self, channel):
        pass

    def on_link(self, beacon):
        pass

    def on_authentication(self, beacon):
        pass

    def on_transport(self, beacon):
        pass

    # Shouldn't have to touch these:

    def start(self):
        self._main()

    def stop(self):
        self._node.stop()

    def erase(self, index):
        pass

    def upload(self, index, data):
        pass

    def download(self, index, callback=None):
        offset  = 0
        initial = True
        crc     = 0
        data    = array.array('B')
        while True:
            _logger.debug("Download %d, o%d, c%d", index, offset, crc)
            self._send_command(DownloadRequest(index, offset, True, crc))
            _logger.debug("Wait for response...")
            try:
                response = self._get_command()
                if response._get_argument("response") == DownloadResponse.Response.OK:
                    remaining    = response._get_argument("remaining")
                    offset       = response._get_argument("offset")
                    total        = offset + remaining
                    data[offset:total] = response._get_argument("data")[:remaining]
                    #print "rem", remaining, "offset", offset, "total", total, "size", response._get_argument("size")
                    if callback != None:
                        callback(float(total) / float(response._get_argument("size")))
                    if total == response._get_argument("size"):
                        return data
                    crc = response._get_argument("crc")
                    offset = total
                else:
                    raise AntFSDownloadException(response._get_argument("response"))
            except Queue.Empty:
                _logger.debug("Download %d timeout", index)
                #print "recover from download failure"

    def download_directory(self, callback=None):
        data = self.download(0, callback)
        return Directory.parse(data)

    def link(self):
        self._channel.request_message(Message.ID.RESPONSE_CHANNEL_ID)
        self._send_command(LinkCommand(self._frequency, 4, self._serial_number))

        # New period, search timeout
        self._channel.set_period(4096)
        self._channel.set_search_timeout(3)
        self._channel.set_rf_freq(self._frequency)

    def authentication_serial(self):
        self._send_command(AuthenticateCommand(
                AuthenticateCommand.Request.SERIAL,
                self._serial_number))
        response = self._get_command()
        return (response.get_serial(), response.get_data_string())

    def authentication_passkey(self, passkey):
        self._send_command(AuthenticateCommand(
                AuthenticateCommand.Request.PASSKEY_EXCHANGE,
                self._serial_number, passkey))

        response = self._get_command()
        if response._get_argument("type") == AuthenticateResponse.Response.ACCEPT:
            return response.get_data_array()
        else:
            raise AntFSAuthenticationException(response._get_argument("type"))

    def authentication_pair(self, friendly_name):
        data = array.array('B', map(ord, list(friendly_name)))
        self._send_command(AuthenticateCommand(
                AuthenticateCommand.Request.PAIRING,
                self._serial_number, data))

        response = self._get_command(30)
        if response._get_argument("type") == AuthenticateResponse.Response.ACCEPT:
            return response.get_data_array()
        else:
            raise AntFSAuthenticationException(response._get_argument("type"))

    def disconnect(self):
        d = DisconnectCommand(DisconnectCommand.Type.RETURN_LINK, 0, 0)
        self._send_command(d)

<|endoftext|>


<commit_before># Copyright 2014 OpenStack Foundation
#
# 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 osc_lib import utils

DEFAULT_ALARMING_API_VERSION = '2'
API_VERSION_OPTION = 'os_alarming_api_version'
API_NAME = "alarming"
API_VERSIONS = {
    "2": "aodhclient.v2.client.Client",
}

def make_client(instance):
    """Returns an queues service client."""
    version = instance._api_version[API_NAME]
    try:
        version = int(version)
    except ValueError:
        version = float(version)

    aodh_client = utils.get_client_class(
        API_NAME,
        version,
        API_VERSIONS)
    # NOTE(sileht): ensure setup of the session is done
    instance.setup_auth()
    return aodh_client(session=instance.session)

def build_option_parser(parser):
    """Hook to add global options."""
    parser.add_argument(
        '--os-alarming-api-version',
        metavar='<alarming-api-version>',
        default=utils.env(
            'OS_ALARMING_API_VERSION',
            default=DEFAULT_ALARMING_API_VERSION),
        help=('Queues API version, default=' +
              DEFAULT_ALARMING_API_VERSION +
              ' (Env: OS_ALARMING_API_VERSION)'))
    return parser
<commit_msg>Add plug-in summary for osc doc<commit_after># Copyright 2014 OpenStack Foundation
#
# 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.

"""OpenStackClient plugin for Telemetry Alarming service."""

from osc_lib import utils

DEFAULT_ALARMING_API_VERSION = '2'
API_VERSION_OPTION = 'os_alarming_api_version'
API_NAME = "alarming"
API_VERSIONS = {
    "2": "aodhclient.v2.client.Client",
}

def make_client(instance):
    """Returns an queues service client."""
    version = instance._api_version[API_NAME]
    try:
        version = int(version)
    except ValueError:
        version = float(version)

    aodh_client = utils.get_client_class(
        API_NAME,
        version,
        API_VERSIONS)
    # NOTE(sileht): ensure setup of the session is done
    instance.setup_auth()
    return aodh_client(session=instance.session)

def build_option_parser(parser):
    """Hook to add global options."""
    parser.add_argument(
        '--os-alarming-api-version',
        metavar='<alarming-api-version>',
        default=utils.env(
            'OS_ALARMING_API_VERSION',
            default=DEFAULT_ALARMING_API_VERSION),
        help=('Queues API version, default=' +
              DEFAULT_ALARMING_API_VERSION +
              ' (Env: OS_ALARMING_API_VERSION)'))
    return parser
<|endoftext|>


<commit_before>"""
    sqmpy.job.saga_helper
    ~~~~~

    Provides ways to interact with saga classes
"""
import os
import time
import base64
import datetime
import hashlib

import saga
import flask
import gevent
from flask.ext.login import current_user

import helpers
from ..database import db
from .constants import FileRelation, ScriptType, HPCBackend
from .exceptions import JobManagerException
from .models import StagingFile, JobStateHistory, Job
from .callback import JobStateChangeCallback

__author__ = 'Mehdi Sadeghi'

class SagaJobWrapper(object):
    """
    To wrap, initialize and run a saga job.
    """
    def __init__(self, job):
        """
        Init
        :param job: an instance of job model class
        :return:
        """
        self._job_id = job.id
        self._job = job

        # If job has resource_job_id, connect to it
        if job.remote_job_id:
            self._job_service = self.make_job_service(job.resource_endpoint)
            self._saga_job = self._job_service.get_job(job.remote_job_id)
        else:
            # Creating the job service object i.e. the connection
            job.resource_endpoint = \
                get_resource_endpoint(job.resource.url, job.hpc_backend)
            self._job_service = self.make_job_service(job.resource_endpoint)

    def make_job_service(self, endpoint):
        # Create ssh security context
        ctx = None
        if flask.current_app.config.get('SSH_WITH_LOGIN_INFO'):
            ctx = saga.Context('userpass')
            ctx.user_id = current_user.username
            ctx.user_pass = base64.b64decode(flask.session['password'])
        else:
            ctx = saga.Context('ssh')
        session = saga.Session()
        # For a reason beyon my current understanding, js instance contains
        # previous contex instances, i.e. it is not fresh. I have to manually
        # delete previous contexts to avoid problems for now.
        # TODO: Fix this
        try:
            while True:
                session.contexts.pop()
        except IndexError:
            pass

        # Explicitely add the only desired security context
        session.add_context(ctx)
        js = saga.job.Service(endpoint, session=session)
        # TODO: Fix in upstream. Service does not populate adaptor's session.
        js._adaptor._set_session(session)
        return js

    def _register_callbacks(self):
        """
        Register callback functions for the saga job
        :return:
        """
        # This callback will locally store output files and new states in db
        self._saga_job.add_callback(saga.STATE,
                                    JobStateChangeCallback(self._job, self))

    def get_job(self):
        """
        Returns the inner job
        :return:
        """
        return self._job

    def run(self):
        """
        Run the job on remote resource
        :return:
        """
        # Set remote job working directory
        remote_job_dir = \
            get_job_endpoint(self._job.id, self._job_service.session)

        # Make sure the working directory is empty
        if remote_job_dir.list():
            raise JobManagerException('Remote directory is not empty')

        # transfer job files to remote directory
        transfer_job_files(self._job.id, remote_job_dir.get_url())

        # Create saga job description
        jd = make_job_description(self._job, remote_job_dir)

        # Create saga job
        self._saga_job = self._job_service.create_job(jd)

        # Register call backs
        # self._register_callbacks()

        # TODO: My monitoring approach is wrong and should be changed.
        # Prepare our gevent greenlet
        @flask.copy_current_request_context
        def monitor_state():
            while True:
                if __debug__:
                    print('Monitoring job %s for changes...' % self._job.id)
                gevent.sleep(3)
                try:
                    val = self._saga_job.state
                    if val != self._job.last_status:
                        # Shout out load
                        helpers.send_state_change_email(self._job.id,
                                                        self._job.owner_id,
                                                        self._job.last_status,
                                                        val,
                                                        silent=True)

                        # Insert history record
                        history_record = JobStateHistory()
                        history_record.change_time = datetime.datetime.now()
                        history_record.old_state = self._job.last_status
                        history_record.new_state = val
                        history_record.job_id = self._job.id
                        db.session.add(history_record)
                        db.session.flush()

                        # If there are new files, transfer them back, along
                        # with output and error files
                        download_job_files(self._job.id,
                                           self._saga_job.description,
                                           self._job_service.session)

                        # Update last status
                        self._job.last_status = val
                        if self._job not in db.session:
                            db.session.merge(self._job)
                        db.session.commit()
                    if val in (saga.FAILED,
                               saga.DONE,
                               saga.CANCELED,
                               saga.FINAL,
                               saga.EXCEPTION):
                        return
                except saga.IncorrectState:
                    pass

        # Spawn the coroutine
        gevent.spawn(monitor_state)

        # Run the job eventually
        flask.current_app.logger.debug("...starting job[%s]..." % self._job.id)
        self._saga_job.run()

        # Store remote pid
        self._job.remote_job_id = self._saga_job.get_id()
        db.session.commit()

        flask.current_app.logger.debug(
            "Remote Job ID    : %s" % self._saga_job.id)
        flask.current_app.logger.debug(
            "Remote Job State : %s" % self._saga_job.state)

    def cancel(self):
        """
        Cancel the job
        """
        self._saga_job.cancel()

def get_resource_endpoint(host, hpc_backend):
    """
    Get ssh URI of remote host
    :param host: host to make url for it
    :param hpc_backend: hpc_backend integer value according to HPCBackend enum
    :return:
    """
    # Default SAGA adaptor to ssh
    adaptor = 'ssh'
    if helpers.is_localhost(host):
        adaptor = 'fork'
    elif hpc_backend == HPCBackend.sge.value:
        adaptor = 'sge+ssh'
    return '{adaptor}://{remote_host}'.format(adaptor=adaptor,
                                              remote_host=host)

def get_job_endpoint(job_id, session):
    """
    Returns the remote job working directory. Creates the parent
    folders if they don't exist.
    :param job_id: job id
    :param session: saga session to be used
    :return:
    """
    job = Job.query.get(job_id)
    # Remote working directory will be inside temp folder
    if not job.remote_dir:
        job.remote_dir = '/tmp/sqmpy/{job_id}'.format(job_id=job.id)
    adapter = 'sftp'
    if helpers.is_localhost(job.resource.url):
        adapter = 'file'
    adaptor_string = '{adapter}://{remote_host}/{working_directory}'
    remote_address = \
        adaptor_string.format(adapter=adapter,
                              remote_host=job.resource.url,
                              working_directory=job.remote_dir)
    # Appropriate folders will be created
    return \
        saga.filesystem.Directory(remote_address,
                                  saga.filesystem.CREATE_PARENTS,
                                  session=session)

def make_job_description(job, remote_job_dir):
    """
    Creates saga job description
    :param job: job instance
    :param remote_job_dir: saga remote job directory instance
    :return:
    """
    script_file = \
        StagingFile.query.filter(
            StagingFile.parent_id == job.id,
            StagingFile.relation == FileRelation.script.value).first()

    jd = saga.job.Description()
    # TODO: Add queue name, project and other params
    jd.working_directory = remote_job_dir.get_url().path
    jd.total_cpu_count = job.total_cpu_count
    jd.wall_time_limit = job.walltime_limit
    jd.spmd_variation = job.spmd_variation
    jd.queue = job.queue or None
    jd.total_physical_memory = job.total_physical_memory or None
    jd.project = job.project or None

    # TODO: Add proper arguments for each input file
    jd.arguments = [script_file.name]

    # TODO: Use script handler instead, see issue #13 on github
    if job.script_type == ScriptType.python.value:
        jd.executable = '/usr/bin/python'
    if job.script_type == ScriptType.shell.value:
        jd.executable = '/bin/sh'
    script_abs_path = '{dir}/{file}'.format(dir=remote_job_dir.get_url().path,
                                            file=script_file.name)
    jd.arguments = [script_abs_path]

    jd.output = '{script_name}.out.txt'.format(script_name=script_file.name)
    jd.error = '{script_name}.err.txt'.format(script_name=script_file.name)
    return jd

def download_job_files(job_id, job_description, session, config=None,
                       wipe=True):
    """
    Copies output and error files along with any other output files back to the
    current machine.
    :param job_id: job id
    :param job_description:
    :param session: saga session to remote resource
    :param wipe: if set to True will wipe files from remote machine.
    :return:
    """
    # Get a new object in this session
    job = Job.query.get(job_id)

    # Get or create job directory
    local_job_dir = helpers.get_job_staging_folder(job_id, config)
    if helpers.is_localhost(job.resource.url):
        local_job_dir_url = local_job_dir
    else:
        local_job_dir_url = \
            helpers.get_job_staging_folder(job_id, config, make_sftp_url=True)

    # Get staging file names for this job which are already uploaded
    # we don't need to download them since we have them already
    uploaded_files = \
        StagingFile.query.with_entities(StagingFile.name).filter(
            StagingFile.parent_id == Job.id,
            Job.id == job_id).all()
    # Convert tuple result to list
    uploaded_files = [file_name for file_name, in uploaded_files]
    remote_dir = get_job_endpoint(job_id, session)
    files = remote_dir.list()
    staging_files = []
    for file_url in files:
        if file_url.path == job_description.output:
            staging_files.append((file_url, FileRelation.output.value))
        elif file_url.path == job_description.error:
            staging_files.append((file_url, FileRelation.error.value))
        elif file_url.path not in uploaded_files:
            staging_files.append((file_url, FileRelation.output.value))

    for file_url, relation in staging_files:
        # Copy physical file to local directory
        if wipe:
            remote_dir.move(file_url, local_job_dir_url)
        else:
            remote_dir.copy(file_url, local_job_dir_url)
        time.sleep(.5)
        # Insert appropriate record into db
        absolute_name = os.path.join(local_job_dir, file_url.path)
        sf = StagingFile()
        sf.name = file_url.path
        sf.relation = relation
        sf.original_name = file_url.path
        sf.checksum = hashlib.md5(open(absolute_name).read()).hexdigest()
        sf.location = local_job_dir
        sf.parent_id = job_id
        db.session.add(sf)
    db.session.commit()

def transfer_job_files(job_id, remote_job_dir_url):
    """
    Upload job files to remote resource
    :param job_id: job id
    :param remote_job_dir_url: saga.url.Url of remote job directory
    :return
    """
    # Copy script and input files to remote host
    uploading_files = \
        StagingFile.query.filter(
            StagingFile.parent_id == job_id,
            StagingFile.relation.in_([FileRelation.input.value,
                                      FileRelation.script.value])).all()
    for file_to_upload in uploading_files:
        file_wrapper = \
            saga.filesystem.File('file://localhost/{file_path}'
                                 .format(file_path=file_to_upload.get_path()))
        file_wrapper.copy(remote_job_dir_url)
<commit_msg>Remove contexts for userpass case<commit_after>"""
    sqmpy.job.saga_helper
    ~~~~~

    Provides ways to interact with saga classes
"""
import os
import time
import base64
import datetime
import hashlib

import saga
import flask
import gevent
from flask.ext.login import current_user

import helpers
from ..database import db
from .constants import FileRelation, ScriptType, HPCBackend
from .exceptions import JobManagerException
from .models import StagingFile, JobStateHistory, Job
from .callback import JobStateChangeCallback

__author__ = 'Mehdi Sadeghi'

class SagaJobWrapper(object):
    """
    To wrap, initialize and run a saga job.
    """
    def __init__(self, job):
        """
        Init
        :param job: an instance of job model class
        :return:
        """
        self._job_id = job.id
        self._job = job

        # If job has resource_job_id, connect to it
        if job.remote_job_id:
            self._job_service = self.make_job_service(job.resource_endpoint)
            self._saga_job = self._job_service.get_job(job.remote_job_id)
        else:
            # Creating the job service object i.e. the connection
            job.resource_endpoint = \
                get_resource_endpoint(job.resource.url, job.hpc_backend)
            self._job_service = self.make_job_service(job.resource_endpoint)

    def make_job_service(self, endpoint):
        # Create ssh security context
        ctx = None
        session = saga.Session()
        if flask.current_app.config.get('SSH_WITH_LOGIN_INFO'):
            ctx = saga.Context('userpass')
            ctx.user_id = current_user.username
            ctx.user_pass = base64.b64decode(flask.session['password'])
            # For a reason beyon my current understanding, js instance contains
            # previous contex instances, i.e. it is not fresh. I have to
            # manually delete previous contexts to avoid problems for now.
            # Moreover this breaks userpass context and renders it useless.
            # TODO: Fix this
            try:
                while True:
                    session.contexts.pop()
            except IndexError:
                pass
        else:
            ctx = saga.Context('ssh')

        # Explicitely add the only desired security context
        session.add_context(ctx)
        js = saga.job.Service(endpoint, session=session)
        # TODO: Fix in upstream. Service does not populate adaptor's session.
        js._adaptor._set_session(session)
        return js

    def _register_callbacks(self):
        """
        Register callback functions for the saga job
        :return:
        """
        # This callback will locally store output files and new states in db
        self._saga_job.add_callback(saga.STATE,
                                    JobStateChangeCallback(self._job, self))

    def get_job(self):
        """
        Returns the inner job
        :return:
        """
        return self._job

    def run(self):
        """
        Run the job on remote resource
        :return:
        """
        # Set remote job working directory
        remote_job_dir = \
            get_job_endpoint(self._job.id, self._job_service.session)

        # Make sure the working directory is empty
        if remote_job_dir.list():
            raise JobManagerException('Remote directory is not empty')

        # transfer job files to remote directory
        transfer_job_files(self._job.id, remote_job_dir.get_url())

        # Create saga job description
        jd = make_job_description(self._job, remote_job_dir)

        # Create saga job
        self._saga_job = self._job_service.create_job(jd)

        # Register call backs
        # self._register_callbacks()

        # TODO: My monitoring approach is wrong and should be changed.
        # Prepare our gevent greenlet
        @flask.copy_current_request_context
        def monitor_state():
            while True:
                if __debug__:
                    print('Monitoring job %s for changes...' % self._job.id)
                gevent.sleep(3)
                try:
                    val = self._saga_job.state
                    if val != self._job.last_status:
                        # Shout out load
                        helpers.send_state_change_email(self._job.id,
                                                        self._job.owner_id,
                                                        self._job.last_status,
                                                        val,
                                                        silent=True)

                        # Insert history record
                        history_record = JobStateHistory()
                        history_record.change_time = datetime.datetime.now()
                        history_record.old_state = self._job.last_status
                        history_record.new_state = val
                        history_record.job_id = self._job.id
                        db.session.add(history_record)
                        db.session.flush()

                        # If there are new files, transfer them back, along
                        # with output and error files
                        download_job_files(self._job.id,
                                           self._saga_job.description,
                                           self._job_service.session)

                        # Update last status
                        self._job.last_status = val
                        if self._job not in db.session:
                            db.session.merge(self._job)
                        db.session.commit()
                    if val in (saga.FAILED,
                               saga.DONE,
                               saga.CANCELED,
                               saga.FINAL,
                               saga.EXCEPTION):
                        return
                except saga.IncorrectState:
                    pass

        # Spawn the coroutine
        gevent.spawn(monitor_state)

        # Run the job eventually
        flask.current_app.logger.debug("...starting job[%s]..." % self._job.id)
        self._saga_job.run()

        # Store remote pid
        self._job.remote_job_id = self._saga_job.get_id()
        db.session.commit()

        flask.current_app.logger.debug(
            "Remote Job ID    : %s" % self._saga_job.id)
        flask.current_app.logger.debug(
            "Remote Job State : %s" % self._saga_job.state)

    def cancel(self):
        """
        Cancel the job
        """
        self._saga_job.cancel()

def get_resource_endpoint(host, hpc_backend):
    """
    Get ssh URI of remote host
    :param host: host to make url for it
    :param hpc_backend: hpc_backend integer value according to HPCBackend enum
    :return:
    """
    # Default SAGA adaptor to ssh
    adaptor = 'ssh'
    if helpers.is_localhost(host):
        adaptor = 'fork'
    elif hpc_backend == HPCBackend.sge.value:
        adaptor = 'sge+ssh'
    return '{adaptor}://{remote_host}'.format(adaptor=adaptor,
                                              remote_host=host)

def get_job_endpoint(job_id, session):
    """
    Returns the remote job working directory. Creates the parent
    folders if they don't exist.
    :param job_id: job id
    :param session: saga session to be used
    :return:
    """
    job = Job.query.get(job_id)
    # Remote working directory will be inside temp folder
    if not job.remote_dir:
        job.remote_dir = '/tmp/sqmpy/{job_id}'.format(job_id=job.id)
    adapter = 'sftp'
    if helpers.is_localhost(job.resource.url):
        adapter = 'file'
    adaptor_string = '{adapter}://{remote_host}/{working_directory}'
    remote_address = \
        adaptor_string.format(adapter=adapter,
                              remote_host=job.resource.url,
                              working_directory=job.remote_dir)
    # Appropriate folders will be created
    return \
        saga.filesystem.Directory(remote_address,
                                  saga.filesystem.CREATE_PARENTS,
                                  session=session)

def make_job_description(job, remote_job_dir):
    """
    Creates saga job description
    :param job: job instance
    :param remote_job_dir: saga remote job directory instance
    :return:
    """
    script_file = \
        StagingFile.query.filter(
            StagingFile.parent_id == job.id,
            StagingFile.relation == FileRelation.script.value).first()

    jd = saga.job.Description()
    # TODO: Add queue name, project and other params
    jd.working_directory = remote_job_dir.get_url().path
    jd.total_cpu_count = job.total_cpu_count
    jd.wall_time_limit = job.walltime_limit
    jd.spmd_variation = job.spmd_variation
    jd.queue = job.queue or None
    jd.total_physical_memory = job.total_physical_memory or None
    jd.project = job.project or None

    # TODO: Add proper arguments for each input file
    jd.arguments = [script_file.name]

    # TODO: Use script handler instead, see issue #13 on github
    if job.script_type == ScriptType.python.value:
        jd.executable = '/usr/bin/python'
    if job.script_type == ScriptType.shell.value:
        jd.executable = '/bin/sh'
    script_abs_path = '{dir}/{file}'.format(dir=remote_job_dir.get_url().path,
                                            file=script_file.name)
    jd.arguments = [script_abs_path]

    jd.output = '{script_name}.out.txt'.format(script_name=script_file.name)
    jd.error = '{script_name}.err.txt'.format(script_name=script_file.name)
    return jd

def download_job_files(job_id, job_description, session, config=None,
                       wipe=True):
    """
    Copies output and error files along with any other output files back to the
    current machine.
    :param job_id: job id
    :param job_description:
    :param session: saga session to remote resource
    :param wipe: if set to True will wipe files from remote machine.
    :return:
    """
    # Get a new object in this session
    job = Job.query.get(job_id)

    # Get or create job directory
    local_job_dir = helpers.get_job_staging_folder(job_id, config)
    if helpers.is_localhost(job.resource.url):
        local_job_dir_url = local_job_dir
    else:
        local_job_dir_url = \
            helpers.get_job_staging_folder(job_id, config, make_sftp_url=True)

    # Get staging file names for this job which are already uploaded
    # we don't need to download them since we have them already
    uploaded_files = \
        StagingFile.query.with_entities(StagingFile.name).filter(
            StagingFile.parent_id == Job.id,
            Job.id == job_id).all()
    # Convert tuple result to list
    uploaded_files = [file_name for file_name, in uploaded_files]
    remote_dir = get_job_endpoint(job_id, session)
    files = remote_dir.list()
    staging_files = []
    for file_url in files:
        if file_url.path == job_description.output:
            staging_files.append((file_url, FileRelation.output.value))
        elif file_url.path == job_description.error:
            staging_files.append((file_url, FileRelation.error.value))
        elif file_url.path not in uploaded_files:
            staging_files.append((file_url, FileRelation.output.value))

    for file_url, relation in staging_files:
        # Copy physical file to local directory
        if wipe:
            remote_dir.move(file_url, local_job_dir_url)
        else:
            remote_dir.copy(file_url, local_job_dir_url)
        time.sleep(.5)
        # Insert appropriate record into db
        absolute_name = os.path.join(local_job_dir, file_url.path)
        sf = StagingFile()
        sf.name = file_url.path
        sf.relation = relation
        sf.original_name = file_url.path
        sf.checksum = hashlib.md5(open(absolute_name).read()).hexdigest()
        sf.location = local_job_dir
        sf.parent_id = job_id
        db.session.add(sf)
    db.session.commit()

def transfer_job_files(job_id, remote_job_dir_url):
    """
    Upload job files to remote resource
    :param job_id: job id
    :param remote_job_dir_url: saga.url.Url of remote job directory
    :return
    """
    # Copy script and input files to remote host
    uploading_files = \
        StagingFile.query.filter(
            StagingFile.parent_id == job_id,
            StagingFile.relation.in_([FileRelation.input.value,
                                      FileRelation.script.value])).all()
    for file_to_upload in uploading_files:
        file_wrapper = \
            saga.filesystem.File('file://localhost/{file_path}'
                                 .format(file_path=file_to_upload.get_path()))
        file_wrapper.copy(remote_job_dir_url)
<|endoftext|>


<commit_before>from collections import OrderedDict

from squad.core.utils import join_name, parse_name
from squad.core import models

class TestComparison(object):
    """
    Data structure:

    builds: [Build]
    environments: Build → [EnvironmentName(str)]
    results: TestName(str) → ((Build,EnvironmentName(str)) → TestResults)

    The best way to think about this is to think of the table you want as
    result:

    +-------+---------------+---------------+
    |       |    build 1    |    build 2    |
    +       +-------+-------+-------+-------+
    |       | env 1 | env 2 | env 1 | env 2 |
    +-------+-------+-------+-------+-------+
    | test1 |       |       |       |       |
    | test2 |       |       |       |       |
    | test3 |       |       |       |       |
    | test4 |       |       |       |       |
    +-------+-------+-------+-------+-------+

    `builds` is the list of builds (the top header row)

    `environments` is a mapping between a build and the list of environments
    (the bottom header row)

    `results` contains the body of the table. It is a mapping between a test
    name string (row) to the results for that row. results[testname] is then a
    mapping, between Environment (the column) and the test result (the cells in
    the table). So results[testname][env] gives you the value of the cell at
    (testname, env)
    """

    def __init__(self, *builds):
        self.builds = list(builds)
        self.environments = OrderedDict()
        self.all_environments = set()
        self.results = OrderedDict()
        self.__intermittent__ = {}

        models.Build.prefetch_related(self.builds)
        self.__extract_results__()

    @classmethod
    def compare_builds(cls, *builds):
        builds = [b for b in builds if b]
        return cls(*builds)

    @classmethod
    def compare_projects(cls, *projects):
        builds = [p.builds.last() for p in projects]
        return cls.compare_builds(*builds)

    def __extract_results__(self):
        for test in self.__all_tests__():
            self.results[test] = OrderedDict()
        for build in self.builds:
            test_runs = list(build.test_runs.all())
            environments = [t.environment for t in test_runs]
            for e in environments:
                self.all_environments.add(str(e))
            self.environments[build] = sorted([str(e) for e in set(environments)])
            for test_run in test_runs:
                self.__extract_test_results__(test_run)

    def __extract_test_results__(self, test_run):
        for test in test_run.tests.prefetch_related('suite', 'known_issues').all():
            key = (test_run.build, str(test_run.environment))
            self.results[test.full_name][key] = test.status
            for issue in test.known_issues.all():
                if issue.intermittent:
                    env = str(test_run.environment)
                    self.__intermittent__[(test.full_name, env)] = True

    def __all_tests__(self):
        test_set = set()
        for build in self.builds:
            for testrun in build.test_runs.all():
                for test in testrun.tests.all():
                    test_set.add(test.full_name)
        return sorted(test_set)

    __diff__ = None

    @property
    def diff(self):
        """
        Returns a subset of the rows, containing only the rows where results
        differ between the builds.
        """
        if self.__diff__ is not None:
            return self.__diff__

        d = OrderedDict()
        for test, results in self.results.items():
            previous = None
            for build in self.builds:
                current = [results.get((build, e)) for e in self.environments[build]]
                if previous and previous != current:
                    d[test] = results
                    break
                previous = current

        self.__diff__ = d
        return self.__diff__

    __regressions__ = None
    __fixes__ = None

    @property
    def regressions(self):
        if self.__regressions__ is None:
            self.__regressions__ = self.__status_changes__(('pass', 'fail'))
        return self.__regressions__

    @property
    def fixes(self):
        if self.__fixes__ is None:
            self.__fixes__ = self.__status_changes__(
                ('fail', 'pass'),
                ('xfail', 'pass'),
                predicate=lambda test, env: (test, env) not in self.__intermittent__
            )
        return self.__fixes__

    def __status_changes__(self, *transitions, predicate=lambda test, env: True):
        if len(self.builds) < 2:
            return {}

        comparisons = OrderedDict()
        after = self.builds[-1]  # last
        before = self.builds[-2]  # second to last
        for env in self.environments[after]:
            comparison_list = []
            for test, results in self.diff.items():
                results_after = results.get((after, env))
                results_before = results.get((before, env))
                if (results_before, results_after) in transitions:
                    if predicate(test, env):
                        comparison_list.append(test)
            if comparison_list:
                comparisons[env] = comparison_list

        return comparisons

    @property
    def regressions_grouped_by_suite(self):
        return self.__status_changes_by_suite__()

    @property
    def fixes_grouped_by_suite(self):
        return self.__status_changes_by_suite__(False)

    def __status_changes_by_suite__(self, regression=True):
        comparisons = self.regressions
        if not regression:
            comparisons = self.fixes
        result = OrderedDict()
        for env, tests in comparisons.items():
            this_env = OrderedDict()
            for test in tests:
                suite, testname = parse_name(test)
                if suite not in this_env:
                    this_env[suite] = []
                this_env[suite].append(testname)
            result[env] = this_env
        return result
<commit_msg>core: TestComparison: optimize database access<commit_after>from collections import OrderedDict

from squad.core.utils import join_name, parse_name
from squad.core import models

class TestComparison(object):
    """
    Data structure:

    builds: [Build]
    environments: Build → [EnvironmentName(str)]
    results: TestName(str) → ((Build,EnvironmentName(str)) → TestResults)

    The best way to think about this is to think of the table you want as
    result:

    +-------+---------------+---------------+
    |       |    build 1    |    build 2    |
    +       +-------+-------+-------+-------+
    |       | env 1 | env 2 | env 1 | env 2 |
    +-------+-------+-------+-------+-------+
    | test1 |       |       |       |       |
    | test2 |       |       |       |       |
    | test3 |       |       |       |       |
    | test4 |       |       |       |       |
    +-------+-------+-------+-------+-------+

    `builds` is the list of builds (the top header row)

    `environments` is a mapping between a build and the list of environments
    (the bottom header row)

    `results` contains the body of the table. It is a mapping between a test
    name string (row) to the results for that row. results[testname] is then a
    mapping, between Environment (the column) and the test result (the cells in
    the table). So results[testname][env] gives you the value of the cell at
    (testname, env)
    """

    def __init__(self, *builds):
        self.builds = list(builds)
        self.environments = OrderedDict()
        self.all_environments = set()
        self.results = OrderedDict()
        self.__intermittent__ = {}

        models.Build.prefetch_related(self.builds)
        self.__extract_results__()

    @classmethod
    def compare_builds(cls, *builds):
        builds = [b for b in builds if b]
        return cls(*builds)

    @classmethod
    def compare_projects(cls, *projects):
        builds = [p.builds.last() for p in projects]
        return cls.compare_builds(*builds)

    def __extract_results__(self):
        for test in self.__all_tests__():
            self.results[test] = OrderedDict()
        test_runs = models.TestRun.objects.filter(
            build__in=self.builds,
        ).prefetch_related(
            'build',
            'environment',
            'tests',
            'tests__suite',
            'tests__known_issues',
        )
        for build in self.builds:
            self.environments[build] = set()
        for test_run in test_runs:
            build = test_run.build
            environment = test_run.environment
            self.all_environments.add(str(environment))
            self.environments[build].add(str(environment))
            self.__extract_test_results__(test_run)
        for build in self.builds:
            self.environments[build] = sorted(self.environments[build])

    def __extract_test_results__(self, test_run):
        for test in test_run.tests.all():
            key = (test_run.build, str(test_run.environment))
            self.results[test.full_name][key] = test.status
            for issue in test.known_issues.all():
                if issue.intermittent:
                    env = str(test_run.environment)
                    self.__intermittent__[(test.full_name, env)] = True

    def __all_tests__(self):
        test_set = set()
        for build in self.builds:
            for testrun in build.test_runs.all():
                for test in testrun.tests.all():
                    test_set.add(test.full_name)
        return sorted(test_set)

    __diff__ = None

    @property
    def diff(self):
        """
        Returns a subset of the rows, containing only the rows where results
        differ between the builds.
        """
        if self.__diff__ is not None:
            return self.__diff__

        d = OrderedDict()
        for test, results in self.results.items():
            previous = None
            for build in self.builds:
                current = [results.get((build, e)) for e in self.environments[build]]
                if previous and previous != current:
                    d[test] = results
                    break
                previous = current

        self.__diff__ = d
        return self.__diff__

    __regressions__ = None
    __fixes__ = None

    @property
    def regressions(self):
        if self.__regressions__ is None:
            self.__regressions__ = self.__status_changes__(('pass', 'fail'))
        return self.__regressions__

    @property
    def fixes(self):
        if self.__fixes__ is None:
            self.__fixes__ = self.__status_changes__(
                ('fail', 'pass'),
                ('xfail', 'pass'),
                predicate=lambda test, env: (test, env) not in self.__intermittent__
            )
        return self.__fixes__

    def __status_changes__(self, *transitions, predicate=lambda test, env: True):
        if len(self.builds) < 2:
            return {}

        comparisons = OrderedDict()
        after = self.builds[-1]  # last
        before = self.builds[-2]  # second to last
        for env in self.environments[after]:
            comparison_list = []
            for test, results in self.diff.items():
                results_after = results.get((after, env))
                results_before = results.get((before, env))
                if (results_before, results_after) in transitions:
                    if predicate(test, env):
                        comparison_list.append(test)
            if comparison_list:
                comparisons[env] = comparison_list

        return comparisons

    @property
    def regressions_grouped_by_suite(self):
        return self.__status_changes_by_suite__()

    @property
    def fixes_grouped_by_suite(self):
        return self.__status_changes_by_suite__(False)

    def __status_changes_by_suite__(self, regression=True):
        comparisons = self.regressions
        if not regression:
            comparisons = self.fixes
        result = OrderedDict()
        for env, tests in comparisons.items():
            this_env = OrderedDict()
            for test in tests:
                suite, testname = parse_name(test)
                if suite not in this_env:
                    this_env[suite] = []
                this_env[suite].append(testname)
            result[env] = this_env
        return result
<|endoftext|>


<commit_before>"""Set up MonoDB schema using monogodngine."""
from textwrap import fill
from mongoengine import Document, EmbeddedDocument
from mongoengine import StringField, IntField, ListField, DictField, MapField
from mongoengine import EmbeddedDocumentField, ReferenceField
from sramongo.sra import SraExperiment

class DocumentString(object):
    def __init__(self, document):
        """Represents mongo Document as strings.

        This is a helper class to make a string representation of a mongo
        document.

        Parameters
        ----------
        document: mongoengine.Document
            A subclass of the mongoengine.Document. Assumes that all data is
            stored in an attribute self._data.

        """
        self.string = ''
        self.add_dict(document._data)

    def add_dict(self, dictionary, spacer=0):
        """Iterates over a dictionary and parses strings.

        Parameters
        ----------
        dictionary: dict
            A dictionary with values to be turned into a string.

        spacer: int
            The amount of indention to use. As we go down further into a
            hierarchy indent values further.
        """
        for k, v in dictionary.items():
            if k.startswith('_'):       # skip private keys
                continue
            if isinstance(v, dict):     # if another dict call itself
                self.string += ' ' * spacer + '{}:\n'.format(k)
                self.add_dict(v, spacer=spacer + 4)
            elif isinstance(v, str) | (v is None):
                self.add_element(k, v, spacer=spacer)
            elif isinstance(v, list):
                self.add_list(k, v, spacer=spacer)
            elif isinstance(v, EmbeddedDocument) | isinstance(v, Document):
                self.add_embeded_document(k, v, spacer=spacer)
            else:
                print(k, v, type(v))
                pass

    def add_element(self, key, value, spacer=0):
        """Add key value pair to string for pretty printing."""
        spacer = ' ' * spacer
        if value is not None:
            # wrap really long lines like abstracts
            if len(value) > 100:
                s = spacer + ' ' * (len(key) + 2)
                value = fill(value, width=75).replace('\n', '\n' + s)
        self.string += spacer + '{}: {}\n'.format(key, value)

    def add_list(self, key, value, spacer=0):
        """Formats ids into astring for printing.

        Parameters
        ----------
        ids: list
            List of ids with a proper string representation (i.e. Xref)
        name: str
            Name of the id will be appended as a title (i.e. external_id)
        """
        spacer = ' ' * spacer
        if len(value) > 0:
            self.string += spacer + '{}:\n'.format(key)
            for x in value:
                s = spacer + ' ' * 4
                v = str(x).strip().replace('\n', ', ')
                self.string += s + '{}\n'.format(v)
        else:
            self.string += spacer + '{}: None\n'.format(key)

    def add_embeded_document(self, key, value, spacer=0):
        """Adds embed document."""
        spacer = ' ' * spacer
        self.string += spacer + '{}:\n'.format(key)
        s = spacer + ' ' * 4
        values = str(value).strip().split('\n')
        v = '\n'.join([s + x for x in values]) + '\n'
        self.string += v

# Different Types of Links
class URLLink(EmbeddedDocument):
    label = StringField()
    url = StringField()

    def __str__(self):
        return DocumentString(self).string

class Xref(EmbeddedDocument):
    db = StringField()
    id = StringField()

    def __str__(self):
        return DocumentString(self).string

class XrefLink(EmbeddedDocument):
    label = StringField()
    db = StringField()
    id = StringField()
    meta = {'allow_inheritance': True}

    def __str__(self):
        return DocumentString(self).string

class EntrezLink(XrefLink):
    query = StringField()

    def __str__(self):
        return DocumentString(self).string

class DDBJLink(XrefLink):
    url = StringField()

    def __str__(self):
        return DocumentString(self).string

class ENALink(XrefLink):
    url = StringField()

    def __str__(self):
        return DocumentString(self).string

# Study
class Submission(EmbeddedDocument):
    submission_id = StringField(primary_key=True)
    broker = StringField()

    # Other IDS
    external_id = ListField(EmbeddedDocumentField(Xref), default=list)
    secondary_id = ListField(EmbeddedDocumentField(Xref), default=list)
    submitter_id = ListField(EmbeddedDocumentField(Xref), default=list)

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, **kwargs):
        sraSubmission = sraExperiment.submission
        sraSubmission.update(kwargs)
        submission = cls(**sraSubmission)
        return submission

class Organization(EmbeddedDocument):
    type = StringField()
    abbreviation = StringField()
    name = StringField()
    email = StringField()
    first_name = StringField()
    last_name = StringField()

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, **kwargs):
        sraOrganization = sraExperiment.organization
        sraOrganization.update(kwargs)
        organization = cls(**sraOrganization)
        return organization

class RelatedStudy(XrefLink):
    is_primary = StringField()

class Study(Document):
    # SRA/DRA/ERA
    study_id = StringField(primary_key=True)

    # Look through the external/secondary/submitter for these database xrefs
    # and pull them out for easy access.
    GEO = StringField()
    GEO_Dataset = StringField()
    BioProject = StringField()
    pubmed = StringField()

    # Other IDs
    external_id = ListField(EmbeddedDocumentField(Xref), default=list)
    secondary_id = ListField(EmbeddedDocumentField(Xref), default=list)
    submitter_id = ListField(EmbeddedDocumentField(Xref), default=list)

    # Attributes
    title = StringField()
    type = StringField()
    abstract = StringField()
    center_name = StringField()
    center_project_name = StringField()
    description = StringField()

    # Related study info
    related_studies = ListField(EmbeddedDocumentField(RelatedStudy), default=list)

    # External links
    url_links = ListField(EmbeddedDocumentField(URLLink), default=list)
    xref_links = ListField(EmbeddedDocumentField(XrefLink), default=list)
    entrez_links = ListField(EmbeddedDocumentField(EntrezLink), default=list)
    ddbj_links = ListField(EmbeddedDocumentField(DDBJLink), default=list)
    ena_links = ListField(EmbeddedDocumentField(ENALink), default=list)

    # NOTE: Additional Fields added post creation
    submission = EmbeddedDocumentField(Submission)
    organization = EmbeddedDocumentField(Organization)
    experiments = ListField(StringField(), default=list)

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, **kwargs):
        sraStudy = sraExperiment.study
        sraStudy.update(kwargs)
        study = cls(**sraStudy)
        study.save()
        return study

# Samples
class Sample(Document):
    # SRS/DRS/ERS
    sample_id = StringField(primary_key=True)

    # Look through the external/secondary/submitter for these database xrefs
    # and pull them out for easy access.
    GEO = StringField()
    BioSample = StringField()
    BioProject = StringField()
    pubmed = StringField()

    # Other IDs
    external_id = ListField(EmbeddedDocumentField(Xref), default=list)
    secondary_id = ListField(EmbeddedDocumentField(Xref), default=list)
    submitter_id = ListField(EmbeddedDocumentField(Xref), default=list)

    # Attributes
    title = StringField()
    taxon_id = StringField()
    scientific_name = StringField()
    common_name = StringField()
    individual_name = StringField()
    description = StringField()
    attributes = DictField()

    # External links
    url_links = ListField(EmbeddedDocumentField(URLLink), default=list)
    xref_links = ListField(EmbeddedDocumentField(XrefLink), default=list)
    entrez_links = ListField(EmbeddedDocumentField(EntrezLink), default=list)
    ddbj_links = ListField(EmbeddedDocumentField(DDBJLink), default=list)
    ena_links = ListField(EmbeddedDocumentField(ENALink), default=list)

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, **kwargs):
        sraSample = sraExperiment.sample
        sraSample.update(kwargs)
        sample = cls(**sraSample)
        sample.save()
        return sample

# Experiment
class Pool(EmbeddedDocument):
    sample_id = StringField()
    GEO = StringField()
    BioSample = StringField()

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment):
        if isinstance(sraExperiment, SraExperiment):
            pools = sraExperiment.pool
        else:
            pools = sraExperiment

        poolList = []
        for p in pools:
            pool = cls()
            pool.sample_id = p.get('sample_id', None)
            pool.GEO = p.get('GEO', None)
            pool.BioSample = p.get('BioSample', None)
            poolList.append(pool)

        return poolList

class Experiment(Document):
    # SRX/DRX/ERX
    experiment_id = StringField(primary_key=True)

    # Look through the external/secondary/submitter for these database xrefs
    # and pull them out for easy access.
    GEO = StringField()
    GEO_Dataset = StringField()
    pubmed = StringField()

    # Other IDs
    external_id = ListField(EmbeddedDocumentField(Xref), default=list)
    secondary_id = ListField(EmbeddedDocumentField(Xref), default=list)
    submitter_id = ListField(EmbeddedDocumentField(Xref), default=list)

    # Attributes
    title = StringField()
    study_id = StringField()
    design = StringField()
    library_name = StringField()
    library_strategy = StringField()
    library_source = StringField()
    library_selection = StringField()
    library_layout = StringField()
    library_layout_orientation = StringField()
    library_layout_length = StringField()
    library_layout_sdev = StringField()
    pooling_strategey = StringField()
    library_construction_protocol = StringField()
    platform = StringField()
    instrument_model = StringField()
    attributes = DictField()

    # External links
    url_links = ListField(EmbeddedDocumentField(URLLink), default=list)
    xref_links = ListField(EmbeddedDocumentField(XrefLink), default=list)
    entrez_links = ListField(EmbeddedDocumentField(EntrezLink), default=list)
    ddbj_links = ListField(EmbeddedDocumentField(DDBJLink), default=list)
    ena_links = ListField(EmbeddedDocumentField(ENALink), default=list)

    # NOTE: Additional Fields added post creation
    study = ReferenceField(Study)

    # NOTE: Additional Fields added post creation
    samples = ListField(EmbeddedDocumentField(Pool))

    # NOTE: Additional Fields added post creation
    runs = ListField(StringField(), default=list)

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, **kwargs):
        sraExp = sraExperiment.experiment
        sraExp.update(kwargs)
        experiment = cls(**sraExp)
        experiment.save()
        return experiment

# Run
class TaxRecord(EmbeddedDocument):
    parent = StringField()
    total_count = StringField()
    self_count = StringField()
    tax_id = StringField()
    rank = StringField()

    def __str__(self):
        return DocumentString(self).string

class TaxAnalysis(EmbeddedDocument):
    nspot_analyze = StringField()
    total_spots = StringField()
    mapped_spots = StringField()
    tax_counts = MapField(EmbeddedDocumentField(TaxRecord))

    def __str__(self):
        return DocumentString(self).string

class Run(Document):
    # SRR/DRR/ERR
    run_id = StringField(primary_key=True)

    # Other IDs
    external_id = ListField(EmbeddedDocumentField(Xref), default=list)
    secondary_id = ListField(EmbeddedDocumentField(Xref), default=list)
    submitter_id = ListField(EmbeddedDocumentField(Xref), default=list)

    # Attributes
    experiment_id = StringField()
    samples = ListField(EmbeddedDocumentField(Pool), default=list)
    nspots = StringField()
    nbases = StringField()
    tax_analysis = EmbeddedDocumentField(TaxAnalysis)
    nreads = StringField()

    # if single ended
    read_count = StringField()
    read_len = StringField()

    # if paired ended
    read_count_r1 = StringField()
    read_len_r1 = StringField()

    read_count_r2 = StringField()
    read_len_r2 = StringField()

    # NOTE: Additional Fields added post creation
    experiment = ReferenceField(Experiment)
    release_date = StringField()
    load_date = StringField()
    size_MB = IntField()
    download_path = StringField()

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, runinfo, **kwargs):
        runs = []
        for sraRun in sraExperiment.run:
            pool = sraRun['samples']
            del sraRun['samples']
            run = cls(**sraRun)
            run.samples.extend(Pool.build_from_SraExperiment(pool))
            run.release_date = runinfo.loc[run.run_id, 'ReleaseDate']
            run.load_date = runinfo.loc[run.run_id, 'LoadDate']
            run.size_MB = runinfo.loc[run.run_id, 'size_MB']
            run.download_path = runinfo.loc[run.run_id, 'download_path']
            run.experiment = Experiment.objects(experiment_id=run.experiment_id).first()
            run.save()
            runs.append(run)
        return runs
<commit_msg>Adds additional exception handling.<commit_after>"""Set up MonoDB schema using monogodngine."""
from textwrap import fill
from mongoengine import Document, EmbeddedDocument
from mongoengine import StringField, IntField, ListField, DictField, MapField
from mongoengine import EmbeddedDocumentField, ReferenceField
from mongoengine.errors import ValidationError
from sramongo.sra import SraExperiment
from sramongo.logger import logger

class DocumentString(object):
    def __init__(self, document):
        """Represents mongo Document as strings.

        This is a helper class to make a string representation of a mongo
        document.

        Parameters
        ----------
        document: mongoengine.Document
            A subclass of the mongoengine.Document. Assumes that all data is
            stored in an attribute self._data.

        """
        self.string = ''
        self.add_dict(document._data)

    def add_dict(self, dictionary, spacer=0):
        """Iterates over a dictionary and parses strings.

        Parameters
        ----------
        dictionary: dict
            A dictionary with values to be turned into a string.

        spacer: int
            The amount of indention to use. As we go down further into a
            hierarchy indent values further.
        """
        for k, v in dictionary.items():
            if k.startswith('_'):       # skip private keys
                continue
            if isinstance(v, dict):     # if another dict call itself
                self.string += ' ' * spacer + '{}:\n'.format(k)
                self.add_dict(v, spacer=spacer + 4)
            elif isinstance(v, str) | (v is None):
                self.add_element(k, v, spacer=spacer)
            elif isinstance(v, list):
                self.add_list(k, v, spacer=spacer)
            elif isinstance(v, EmbeddedDocument) | isinstance(v, Document):
                self.add_embeded_document(k, v, spacer=spacer)
            else:
                print(k, v, type(v))
                pass

    def add_element(self, key, value, spacer=0):
        """Add key value pair to string for pretty printing."""
        spacer = ' ' * spacer
        if value is not None:
            # wrap really long lines like abstracts
            if len(value) > 100:
                s = spacer + ' ' * (len(key) + 2)
                value = fill(value, width=75).replace('\n', '\n' + s)
        self.string += spacer + '{}: {}\n'.format(key, value)

    def add_list(self, key, value, spacer=0):
        """Formats ids into astring for printing.

        Parameters
        ----------
        ids: list
            List of ids with a proper string representation (i.e. Xref)
        name: str
            Name of the id will be appended as a title (i.e. external_id)
        """
        spacer = ' ' * spacer
        if len(value) > 0:
            self.string += spacer + '{}:\n'.format(key)
            for x in value:
                s = spacer + ' ' * 4
                v = str(x).strip().replace('\n', ', ')
                self.string += s + '{}\n'.format(v)
        else:
            self.string += spacer + '{}: None\n'.format(key)

    def add_embeded_document(self, key, value, spacer=0):
        """Adds embed document."""
        spacer = ' ' * spacer
        self.string += spacer + '{}:\n'.format(key)
        s = spacer + ' ' * 4
        values = str(value).strip().split('\n')
        v = '\n'.join([s + x for x in values]) + '\n'
        self.string += v

# Different Types of Links
class URLLink(EmbeddedDocument):
    label = StringField()
    url = StringField()

    def __str__(self):
        return DocumentString(self).string

class Xref(EmbeddedDocument):
    db = StringField()
    id = StringField()

    def __str__(self):
        return DocumentString(self).string

class XrefLink(EmbeddedDocument):
    label = StringField()
    db = StringField()
    id = StringField()
    meta = {'allow_inheritance': True}

    def __str__(self):
        return DocumentString(self).string

class EntrezLink(XrefLink):
    query = StringField()

    def __str__(self):
        return DocumentString(self).string

class DDBJLink(XrefLink):
    url = StringField()

    def __str__(self):
        return DocumentString(self).string

class ENALink(XrefLink):
    url = StringField()

    def __str__(self):
        return DocumentString(self).string

# Study
class Submission(EmbeddedDocument):
    submission_id = StringField(primary_key=True)
    broker = StringField()

    # Other IDS
    external_id = ListField(EmbeddedDocumentField(Xref), default=list)
    secondary_id = ListField(EmbeddedDocumentField(Xref), default=list)
    submitter_id = ListField(EmbeddedDocumentField(Xref), default=list)

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, **kwargs):
        try:
            sraSubmission = sraExperiment.submission
            sraSubmission.update(kwargs)
            submission = cls(**sraSubmission)
            return submission
        except ValidationError as err:
            logger.error('ValidationError: submission malformed.')
            logger.debug(sraSubmission)
            return None

class Organization(EmbeddedDocument):
    type = StringField()
    abbreviation = StringField()
    name = StringField()
    email = StringField()
    first_name = StringField()
    last_name = StringField()

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, **kwargs):
        try:
            sraOrganization = sraExperiment.organization
            sraOrganization.update(kwargs)
            organization = cls(**sraOrganization)
            return organization
        except ValidationError as err:
            logger.error('ValidationError: organization malformed.')
            logger.debug(sraOrganization)
            return None

class RelatedStudy(XrefLink):
    is_primary = StringField()

class Study(Document):
    # SRA/DRA/ERA
    study_id = StringField(primary_key=True)

    # Look through the external/secondary/submitter for these database xrefs
    # and pull them out for easy access.
    GEO = StringField()
    GEO_Dataset = StringField()
    BioProject = StringField()
    pubmed = StringField()

    # Other IDs
    external_id = ListField(EmbeddedDocumentField(Xref), default=list)
    secondary_id = ListField(EmbeddedDocumentField(Xref), default=list)
    submitter_id = ListField(EmbeddedDocumentField(Xref), default=list)

    # Attributes
    title = StringField()
    type = StringField()
    abstract = StringField()
    center_name = StringField()
    center_project_name = StringField()
    description = StringField()

    # Related study info
    related_studies = ListField(EmbeddedDocumentField(RelatedStudy), default=list)

    # External links
    url_links = ListField(EmbeddedDocumentField(URLLink), default=list)
    xref_links = ListField(EmbeddedDocumentField(XrefLink), default=list)
    entrez_links = ListField(EmbeddedDocumentField(EntrezLink), default=list)
    ddbj_links = ListField(EmbeddedDocumentField(DDBJLink), default=list)
    ena_links = ListField(EmbeddedDocumentField(ENALink), default=list)

    # NOTE: Additional Fields added post creation
    submission = EmbeddedDocumentField(Submission)
    organization = EmbeddedDocumentField(Organization)
    experiments = ListField(StringField(), default=list)

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, **kwargs):
        try:
            sraStudy = sraExperiment.study
            sraStudy.update(kwargs)
            study = cls(**sraStudy)
            study.save()
            return study
        except ValidationError as err:
            logger.error('ValidationError: study malformed.')
            logger.debug(sraStudy)
            return None

# Samples
class Sample(Document):
    # SRS/DRS/ERS
    sample_id = StringField(primary_key=True)

    # Look through the external/secondary/submitter for these database xrefs
    # and pull them out for easy access.
    GEO = StringField()
    BioSample = StringField()
    BioProject = StringField()
    pubmed = StringField()

    # Other IDs
    external_id = ListField(EmbeddedDocumentField(Xref), default=list)
    secondary_id = ListField(EmbeddedDocumentField(Xref), default=list)
    submitter_id = ListField(EmbeddedDocumentField(Xref), default=list)

    # Attributes
    title = StringField()
    taxon_id = StringField()
    scientific_name = StringField()
    common_name = StringField()
    individual_name = StringField()
    description = StringField()
    attributes = DictField()

    # External links
    url_links = ListField(EmbeddedDocumentField(URLLink), default=list)
    xref_links = ListField(EmbeddedDocumentField(XrefLink), default=list)
    entrez_links = ListField(EmbeddedDocumentField(EntrezLink), default=list)
    ddbj_links = ListField(EmbeddedDocumentField(DDBJLink), default=list)
    ena_links = ListField(EmbeddedDocumentField(ENALink), default=list)

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, **kwargs):
        try:
            sraSample = sraExperiment.sample
            sraSample.update(kwargs)
            sample = cls(**sraSample)
            sample.save()
            return sample
        except ValidationError as err:
            logger.error('ValidationError: sample malformed.')
            logger.debug(sraSample)
            return None

# Experiment
class Pool(EmbeddedDocument):
    sample_id = StringField()
    GEO = StringField()
    BioSample = StringField()

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment):
        if sraExperiment is None:
            return []
        elif isinstance(sraExperiment, SraExperiment):
            pools = sraExperiment.pool
        else:
            pools = sraExperiment

        poolList = []
        for p in pools:
            pool = cls()
            pool.sample_id = p.get('sample_id', None)
            pool.GEO = p.get('GEO', None)
            pool.BioSample = p.get('BioSample', None)
            poolList.append(pool)

        return poolList

class Experiment(Document):
    # SRX/DRX/ERX
    experiment_id = StringField(primary_key=True)

    # Look through the external/secondary/submitter for these database xrefs
    # and pull them out for easy access.
    GEO = StringField()
    GEO_Dataset = StringField()
    pubmed = StringField()

    # Other IDs
    external_id = ListField(EmbeddedDocumentField(Xref), default=list)
    secondary_id = ListField(EmbeddedDocumentField(Xref), default=list)
    submitter_id = ListField(EmbeddedDocumentField(Xref), default=list)

    # Attributes
    title = StringField()
    study_id = StringField()
    design = StringField()
    library_name = StringField()
    library_strategy = StringField()
    library_source = StringField()
    library_selection = StringField()
    library_layout = StringField()
    library_layout_orientation = StringField()
    library_layout_length = StringField()
    library_layout_sdev = StringField()
    pooling_strategey = StringField()
    library_construction_protocol = StringField()
    platform = StringField()
    instrument_model = StringField()
    attributes = DictField()

    # External links
    url_links = ListField(EmbeddedDocumentField(URLLink), default=list)
    xref_links = ListField(EmbeddedDocumentField(XrefLink), default=list)
    entrez_links = ListField(EmbeddedDocumentField(EntrezLink), default=list)
    ddbj_links = ListField(EmbeddedDocumentField(DDBJLink), default=list)
    ena_links = ListField(EmbeddedDocumentField(ENALink), default=list)

    # NOTE: Additional Fields added post creation
    study = ReferenceField(Study)

    # NOTE: Additional Fields added post creation
    samples = ListField(EmbeddedDocumentField(Pool))

    # NOTE: Additional Fields added post creation
    runs = ListField(StringField(), default=list)

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, **kwargs):
        try:
            sraExp = sraExperiment.experiment
            sraExp.update(kwargs)
            experiment = cls(**sraExp)
            experiment.save()
            return experiment
        except ValidationError as err:
            logger.error('ValidationError: experiment malformed.')
            logger.debug(sraExp)
            return None

# Run
class TaxRecord(EmbeddedDocument):
    parent = StringField()
    total_count = StringField()
    self_count = StringField()
    tax_id = StringField()
    rank = StringField()

    def __str__(self):
        return DocumentString(self).string

class TaxAnalysis(EmbeddedDocument):
    nspot_analyze = StringField()
    total_spots = StringField()
    mapped_spots = StringField()
    tax_counts = MapField(EmbeddedDocumentField(TaxRecord))

    def __str__(self):
        return DocumentString(self).string

class Run(Document):
    # SRR/DRR/ERR
    run_id = StringField(primary_key=True)

    # Other IDs
    external_id = ListField(EmbeddedDocumentField(Xref), default=list)
    secondary_id = ListField(EmbeddedDocumentField(Xref), default=list)
    submitter_id = ListField(EmbeddedDocumentField(Xref), default=list)

    # Attributes
    experiment_id = StringField()
    samples = ListField(EmbeddedDocumentField(Pool), default=list)
    nspots = StringField()
    nbases = StringField()
    tax_analysis = EmbeddedDocumentField(TaxAnalysis)
    nreads = StringField()

    # if single ended
    read_count = StringField()
    read_len = StringField()

    # if paired ended
    read_count_r1 = StringField()
    read_len_r1 = StringField()

    read_count_r2 = StringField()
    read_len_r2 = StringField()

    # NOTE: Additional Fields added post creation
    experiment = ReferenceField(Experiment)
    release_date = StringField()
    load_date = StringField()
    size_MB = IntField()
    download_path = StringField()

    def __str__(self):
        return DocumentString(self).string

    @classmethod
    def build_from_SraExperiment(cls, sraExperiment, runinfo, **kwargs):
        runs = []
        for sraRun in sraExperiment.run:
            try:
                if 'samples' in sraRun:
                    pool = sraRun['samples']
                    del sraRun['samples']
                else:
                    pool = None

                run = cls(**sraRun)
                run.samples.extend(Pool.build_from_SraExperiment(pool))

                try:
                    run.release_date = runinfo.loc[run.run_id, 'ReleaseDate']
                except KeyError:
                    pass

                try:
                    run.load_date = runinfo.loc[run.run_id, 'LoadDate']
                except KeyError:
                    pass

                try:
                    run.size_MB = runinfo.loc[run.run_id, 'size_MB']
                except KeyError:
                    pass

                try:
                    run.download_path = runinfo.loc[run.run_id, 'download_path']
                except KeyError:
                    pass

                try:
                    run.experiment = Experiment.objects(experiment_id=run.experiment_id).first()
                except:
                    pass

                run.save()
                runs.append(run)
            except ValidationError as err:
                logger.error('ValidationError: run malformed.')
                logger.debug(sraRun)
                logger.debug(runinfo.loc[run.run_id, :])

        return runs
<|endoftext|>