<commit_before>#!/usr/bin/env python
import time
import threading
import Queue
import os
import pyuarm
import rospy

from std_msgs.msg import String

def ros_try_catch(fn):
        def try_catch(*args, **kwargs):
            try:
                return fn(*args, **kwargs)
            except Exception as e:
                err_msg = "Error in function '" + fn.func_name + "': " + str(e.message)
                rospy.logerr(err_msg)
                rospy.signal_shutdown(err_msg)
        return try_catch

class UARM_interface():

    def __init__(self):
        self.uarm = None
        self.connected = False
        self.alive = True

        self.ros_rate = None
        self.ros_hz = 0

        self.settings = []

        self.uarm_read_thread = None
        self.uarm_read_queue = None
        self.read_queue_get_lock = threading.Lock()
        self.read_queue_put_lock = threading.Lock()
        self.uarm_read_pub = None

        self.uarm_interface_thread = None
        self.uarm_interface_queue = None
        self.interface_queue_get_lock = threading.Lock()
        self.interface_queue_put_lock = threading.Lock()

        self.playback_thread = None
        self.playback_active = False
        self.playback_data = []
        self.loading = False

        self.ja = []

    def get_setting(self, name, string_val=False, list=False):
        for line in self.settings:
            if line[0] == name:
                if string_val:
                    return line[1]
                else:
                    if list:
                        return float(line[1:])
                    else:
                        return float(line[1])

    def connect(self):
        try:
            self.uarm = pyuarm.get_uarm()
            self.connected = True
            self.load_parameters()
            self.start_threads()
            rospy.loginfo("Connected")
        except Exception as e:
            self.connected = False
            err = str(e.message)
            rospy.logerr("Connection error: " + err)
            rospy.signal_shutdown("Connection error")

    def shutdown(self):
        rospy.loginfo("Shutdown function call")
        rospy.signal_shutdown("Shutdown function call")
        start_shutdown = time.time()
        while self.uarm_interface_thread.is_alive() or self.uarm_read_thread.is_alive():
            # self.send_to_interface_queue("SHUTDOWN", priority=True)
            # self.send_to_read_queue("SHUTDOWN", priority=True)
            self.send_to_interface_queue("SHUTDOWN", priority=1)
            self.send_to_read_queue("SHUTDOWN", priority=1)
            time.sleep(0.5)
            if time.time()-start_shutdown > 5:
                break
        self.alive = False
        self.uarm.disconnect()

    def load_parameters(self):
        rospy.loginfo("Loading parameters")

        dir_path = os.path.dirname(os.path.realpath(__file__))
        f = open(dir_path + "/settings.txt")
        self.settings = []
        for line in f:
            self.settings.append(line.split(','))
        f.close()

        self.ros_hz = self.get_setting("ros_hz")
        self.read_pos = self.get_setting("read_pos")
        self.read_ja = self.get_setting("read_ja")
        self.read_AI = self.get_setting("read_AI",list=True)

        self.uarm_read_pub = rospy.Publisher('uarm_read', String, queue_size=10)
        rospy.Subscriber("uarm_write", String, self.uarm_write_callback)
        rospy.init_node('uarm_node', anonymous=True)
        self.ros_rate = rospy.Rate(self.ros_hz)

        #self.uarm_read_queue = Queue.Queue()
        #self.uarm_interface_queue = Queue.Queue()

        self.uarm_read_queue = Queue.PriorityQueue()
        self.uarm_interface_queue = Queue.PriorityQueue()

        rospy.loginfo("Done loading parameters.")

    def start_threads(self):
        rospy.loginfo("Starting threads")
        if self.connected is True:
            self.uarm_read_thread = threading.Thread(target=self.uarm_read)
            self.uarm_interface_thread = threading.Thread(target=self.uarm_interface)

            self.uarm_read_thread.daemon = True
            self.uarm_read_thread.start()

            self.uarm_interface_thread.daemon = True
            self.uarm_interface_thread.start()

            self.request_detach()

        else:
            rospy.logerr("Startup error")
            rospy.signal_shutdown("Startup error")

    def get_read_data(self):
        msg = []
        print "here"
        if self.read_pos == 1:
            msg.append(self.read_position())
        elif self.read_ja == 1:
            msg.append(self.read_joint_angles())
        elif self.read_AI[0] == 1:
            for i in self.read_AI[1:]:
                msg.append(self.read_analog(i))

        print msg
        return msg

    def uarm_interface(self):
        rospy.loginfo("uarm_interface running")
        while True and (rospy.is_shutdown() is False):
            request = self.get_from_interface_queue()
            if request == "READ":
                curr_vals = self.get_read_data()
                try:
                    self.send_to_read_queue(curr_vals)
                except Exception as e:
                    print "Error: ", e
                    rospy.logerr("Error writing to queue. Queue probably shutdown.")
                    break
                #time.sleep(1.0/self.ros_hz)
                try:
                    self.ja = curr_vals[1]
                except Exception as e:
                    err = str(e.message)
                    rospy.logerr("uArm not ready! " + err)
                    rospy.signal_shutdown("Error reading from uArm")

                if self.playback_active is False and self.loading is False:
                    self.send_to_interface_queue("READ")

            elif request == "SHUTDOWN":
                break
            else:
                self.process_command(request)
                if self.playback_active is False and self.loading is False:
                    self.send_to_interface_queue("READ")

        rospy.loginfo("uarm_interface shutdown")
        self.shutdown()

    def request_detach(self):
        rospy.loginfo("uArm detach")
        self.send_to_interface_queue("DET")

    def request_attach(self):
        rospy.loginfo("uArm attach")
        self.send_to_interface_queue("ATT")

    def pump_on(self):
        rospy.loginfo("uArm Pump On")
        self.send_to_interface_queue("PUMP_ON")

    def pump_off(self):
        rospy.loginfo("uArm Pump Off")
        self.send_to_interface_queue("PUMP_OFF")

    @ros_try_catch
    def read_analog(self, pin_num):
        return round(self.uarm.get_analog(pin_num), 3)

    @ros_try_catch
    def get_joint_angles(self):
        return self.uarm.get_servo_angle()

    @ros_try_catch
    def read_position(self):
        return self.uarm.get_position()

    @ros_try_catch
    def read_joint_angles(self):
        # return [self.get_joint_angles(), time.time()]
        return self.get_joint_angles()

    def uarm_read(self):
        rospy.loginfo("uarm_read running")
        while True and (rospy.is_shutdown() is False):
            robot_values = self.get_from_read_queue(blocking=False)
            if robot_values == "SHUTDOWN":
                rospy.logwarn("sending shutdown signal")
                rospy.signal_shutdown("Normal Shutdown Procedure")
                break
            elif not robot_values:
                pass
            elif robot_values == "DONE":
                self.uarm_read_pub.publish("DONE")
            else:
                msg = str(robot_values)
                msg = msg.translate(None, '[]')
                self.uarm_read_pub.publish(msg)
                #self.ros_rate.sleep()

        rospy.loginfo("uarm_read shutdown")

    def send_to_interface_queue(self, msg, msg_list=False, priority=10):
        with self.interface_queue_put_lock:
            if msg_list:
                for m in msg_list:
                    self.uarm_interface_queue.put((priority,m))
            else:
                self.uarm_interface_queue.put((priority, msg))

    # def send_to_interface_queue(self, msg, msg_list=False, priority=False):
    #     with self.interface_queue_put_lock:
    #         if msg_list:
    #             if priority:
    #                 msg_list = msg + self.get_from_interface_queue(all_msgs=True)
    #             for m in msg_list:
    #                 self.uarm_interface_queue.put(m)
    #         else:
    #             if priority:
    #                 msg_list = [msg] + self.get_from_interface_queue(all_msgs=True)
    #                 for m in msg_list:
    #                     self.uarm_interface_queue.put(m)
    #             else:
    #                 self.uarm_interface_queue.put(msg)

    def filter_interface_queue(self, msg_filter, invert=False):
        msgs = self.get_from_interface_queue(all_msgs=True, blocking=True)
        if msg_filter is not "":
            # print "FILTER: ", filter
            n = 0
            for i, msg in enumerate(msgs):
                if filter in msg and not invert:
                    msgs.pop(i)
                    n += 1
                elif filter not in msg and invert:
                    msgs.pop(i)
                    n += 1
            if msgs is not []:
                # print n, " messages filtered matching search terms", filter, invert
                # print msgs
                #self.send_to_interface_queue(msgs, msg_list=True, priority=True)
                self.send_to_interface_queue(msgs, msg_list=True, priority=1)

    def get_from_interface_queue(self, all_msgs=False, blocking=False):
        with self.interface_queue_get_lock:
            if all_msgs:
                msgs = []
                while self.uarm_interface_queue.qsize() > 0:
                    try:
                        msg = self.uarm_interface_queue.get(blocking)
                        try:
                            msg = msg[1]
                        except:
                            pass
                        print msg
                        msgs.append(msg)

                    except Queue.Empty:
                        msgs = []

                return msgs
            else:
                try:
                    msg = self.uarm_interface_queue.get(blocking)
                except Queue.Empty:
                    msg = ""
                try:
                    msg = msg[1]
                except:
                    pass
                return msg

    def send_to_read_queue(self, msg, msg_list=False, priority=10):
        with self.read_queue_put_lock:
            if msg_list:
                for m in msg:
                    self.uarm_read_queue.put((priority,m))
            else:
                self.uarm_read_queue.put((priority,msg))

    # def send_to_read_queue(self, msg, msg_list=False, priority=False):
    #     with self.read_queue_put_lock:
    #         if msg_list:
    #             if priority:
    #                 msg_list = msg + self.get_from_read_queue(all_msgs=True)
    #             for m in msg_list:
    #                 self.uarm_read_queue.put(m)
    #         else:
    #             if priority:
    #                 # print "Sending priority message, read_queue:", msg
    #                 msg_list = [msg] + self.get_from_read_queue(all_msgs=True)
    #                 for m in msg_list:
    #                     self.uarm_read_queue.put(m)
    #             else:
    #                 self.uarm_read_queue.put(msg)

    def get_from_read_queue(self, all_msgs=False, blocking=False):
        with self.read_queue_get_lock:
            if all_msgs:
                msgs = []
                while self.uarm_read_queue.qsize() > 0:
                    try:
                        msg = self.uarm_read_queue.get(blocking)
                        try:
                            msg = msg[1:]
                        except:
                            pass
                        msgs.append()
                    except Queue.Empty:
                        msgs = []
                return msgs
            else:
                try:
                    msg = self.uarm_read_queue.get(blocking)
                    try:
                        msg = msg[1:]
                    except:
                        pass

                except Queue.Empty:
                    msg = ""
                return msg

    def uarm_write_callback(self, data):
        # print "DATA: ", data.data
        if data.data == "CLEAR":
                self.get_from_interface_queue(all_msgs=True)
                self.get_from_read_queue(all_msgs=True)

        elif data.data[0] == "!":
                msg = data.data[1:]
                #print msg
                #self.send_to_interface_queue(msg, priority=True)
                self.send_to_interface_queue(msg, priority=1)
        else:
                self.send_to_interface_queue(data.data)

    def playback(self):
        if self.playback_data:
            self.loading = False
            self.get_from_interface_queue(all_msgs=True)  # flush queue
            try:
                joint_angles = self.ja
                self.send_to_interface_queue("JA"+str(joint_angles[0]) + "," +
                                             str(joint_angles[1]) + "," +
                                             str(joint_angles[2]) +
                                             str(joint_angles[3]))
            except Exception as e:
                # print "joint angles not ready yet: ", e
                err = str(e.message)
                rospy.logerr("joint angles not ready yet: " + err)
                rospy.signal_shutdown("joint angles not ready yet")

            self.send_to_interface_queue("ATT")

            self.filter_interface_queue("READ")

            for line in self.playback_data:
                # print self.playback_active, line
                if self.playback_active is False:
                    self.get_from_interface_queue(all_msgs=True)
                    break
                msg = str(line)

                self.send_to_interface_queue("READ")
                self.send_to_interface_queue(msg)

        self.send_to_interface_queue("BEEP")
        self.send_to_interface_queue("DET")
        self.playback_data = []

    def load_playback_data(self, data_line):
        self.loading = True
        rospy.loginfo("Loading playback data...")
        self.playback_data.append(data_line)

    @ros_try_catch
    def process_command(self, command):
        if command == "PUMP_ON":
            self.uarm.set_pump(1)

        if command == "PUMP_OFF":
            self.uarm.set_pump(0)

        if command == "ATT":
            self.uarm.set_servo_attach()

        if command == "DET":
            self.uarm.set_servo_detach()

        if command[0:3] == "POS" and self.playback_active:
            position = map(float, command[3:].split(','))
            self.uarm.set_position(position[0], position[1], position[2])

        if command[0:3] == "REL":
            adjust = float(command[3:])
            # print adjust
            self.uarm.set_position(x=0, y=0, z=adjust, speed=0, relative=True)

        if command[0:2] == "JA" and self.playback_active:
            angle = map(float, command[2:].split(','))
            for i in range(0, len(angle)-1):
                self.uarm.set_servo_angle(i, angle[i])

        if command[0:4] == "LOAD":
            self.load_playback_data(command[4:])

        if command == "PLAY":
            if self.playback_data:
                self.playback_active = True
                self.playback_thread = threading.Thread(target=self.playback)
                self.playback_thread.daemon = True

                self.playback_thread.start()
        if command == "BEEP":
            self.uarm.set_buzzer(10000, 0.1)
            #self.send_to_read_queue("DONE", priority=True)
            self.send_to_read_queue("DONE", priority=1)
            self.playback_active = False

        if command == "STOP":
            self.playback_active = False
            self.filter_interface_queue(msg_filter="READ")
            time.sleep(0.5)
            self.filter_interface_queue(msg_filter="JA")
            self.playback_data = []
<commit_msg>speed debug<commit_after>#!/usr/bin/env python
import time
import threading
import Queue
import os
import pyuarm
import rospy

from std_msgs.msg import String

def ros_try_catch(fn):
        def try_catch(*args, **kwargs):
            try:
                return fn(*args, **kwargs)
            except Exception as e:
                err_msg = "Error in function '" + fn.func_name + "': " + str(e.message)
                rospy.logerr(err_msg)
                rospy.signal_shutdown(err_msg)
        return try_catch

class UARM_interface():

    def __init__(self):
        self.uarm = None
        self.connected = False
        self.alive = True

        self.ros_rate = None
        self.ros_hz = 0

        self.settings = []

        self.uarm_read_thread = None
        self.uarm_read_queue = None
        self.read_queue_get_lock = threading.Lock()
        self.read_queue_put_lock = threading.Lock()
        self.uarm_read_pub = None

        self.uarm_interface_thread = None
        self.uarm_interface_queue = None
        self.interface_queue_get_lock = threading.Lock()
        self.interface_queue_put_lock = threading.Lock()

        self.playback_thread = None
        self.playback_active = False
        self.playback_data = []
        self.loading = False

        self.ja = []

    def get_setting(self, name, string_val=False, list=False):
        for line in self.settings:
            if line[0] == name:
                if string_val:
                    return line[1]
                else:
                    if list:
                        return float(line[1:])
                    else:
                        return float(line[1])

    def connect(self):
        try:
            self.uarm = pyuarm.get_uarm()
            self.connected = True
            self.load_parameters()
            self.start_threads()
            rospy.loginfo("Connected")
        except Exception as e:
            self.connected = False
            err = str(e.message)
            rospy.logerr("Connection error: " + err)
            rospy.signal_shutdown("Connection error")

    def shutdown(self):
        rospy.loginfo("Shutdown function call")
        rospy.signal_shutdown("Shutdown function call")
        start_shutdown = time.time()
        while self.uarm_interface_thread.is_alive() or self.uarm_read_thread.is_alive():
            # self.send_to_interface_queue("SHUTDOWN", priority=True)
            # self.send_to_read_queue("SHUTDOWN", priority=True)
            self.send_to_interface_queue("SHUTDOWN", priority=1)
            self.send_to_read_queue("SHUTDOWN", priority=1)
            time.sleep(0.5)
            if time.time()-start_shutdown > 5:
                break
        self.alive = False
        self.uarm.disconnect()

    def load_parameters(self):
        rospy.loginfo("Loading parameters")

        dir_path = os.path.dirname(os.path.realpath(__file__))
        f = open(dir_path + "/settings.txt")
        self.settings = []
        for line in f:
            self.settings.append(line.split(','))
        f.close()

        self.ros_hz = self.get_setting("ros_hz")
        self.read_pos = self.get_setting("read_pos")
        self.read_ja = self.get_setting("read_ja")
        self.read_AI = self.get_setting("read_AI",list=True)

        self.uarm_read_pub = rospy.Publisher('uarm_read', String, queue_size=10)
        rospy.Subscriber("uarm_write", String, self.uarm_write_callback)
        rospy.init_node('uarm_node', anonymous=True)
        self.ros_rate = rospy.Rate(self.ros_hz)

        #self.uarm_read_queue = Queue.Queue()
        #self.uarm_interface_queue = Queue.Queue()

        self.uarm_read_queue = Queue.PriorityQueue()
        self.uarm_interface_queue = Queue.PriorityQueue()

        rospy.loginfo("Done loading parameters.")

    def start_threads(self):
        rospy.loginfo("Starting threads")
        if self.connected is True:
            self.uarm_read_thread = threading.Thread(target=self.uarm_read)
            self.uarm_interface_thread = threading.Thread(target=self.uarm_interface)

            self.uarm_read_thread.daemon = True
            self.uarm_read_thread.start()

            self.uarm_interface_thread.daemon = True
            self.uarm_interface_thread.start()

            self.request_detach()

        else:
            rospy.logerr("Startup error")
            rospy.signal_shutdown("Startup error")

    def get_read_data(self):
        msg = []
        msg.append(self.read_analog())
        # print "here"
        # if self.read_pos == 1:
        #     msg.append(self.read_position())
        # elif self.read_ja == 1:
        #     msg.append(self.read_joint_angles())
        # elif self.read_AI[0] == 1:
        #     for i in self.read_AI[1:]:
        #         msg.append(self.read_analog(i))
        #
        # print msg
        return msg

    def uarm_interface(self):
        rospy.loginfo("uarm_interface running")
        while True and (rospy.is_shutdown() is False):
            request = self.get_from_interface_queue()
            if request == "READ":
                curr_vals = self.get_read_data()
                try:
                    self.send_to_read_queue(curr_vals)
                except Exception as e:
                    print "Error: ", e
                    rospy.logerr("Error writing to queue. Queue probably shutdown.")
                    break
                #time.sleep(1.0/self.ros_hz)
                try:
                    self.ja = curr_vals[1]
                except Exception as e:
                    err = str(e.message)
                    rospy.logerr("uArm not ready! " + err)
                    rospy.signal_shutdown("Error reading from uArm")

                if self.playback_active is False and self.loading is False:
                    self.send_to_interface_queue("READ")

            elif request == "SHUTDOWN":
                break
            else:
                self.process_command(request)
                if self.playback_active is False and self.loading is False:
                    self.send_to_interface_queue("READ")

        rospy.loginfo("uarm_interface shutdown")
        self.shutdown()

    def request_detach(self):
        rospy.loginfo("uArm detach")
        self.send_to_interface_queue("DET")

    def request_attach(self):
        rospy.loginfo("uArm attach")
        self.send_to_interface_queue("ATT")

    def pump_on(self):
        rospy.loginfo("uArm Pump On")
        self.send_to_interface_queue("PUMP_ON")

    def pump_off(self):
        rospy.loginfo("uArm Pump Off")
        self.send_to_interface_queue("PUMP_OFF")

    @ros_try_catch
    def read_analog(self, pin_num):
        return round(self.uarm.get_analog(pin_num), 3)

    @ros_try_catch
    def get_joint_angles(self):
        return self.uarm.get_servo_angle()

    @ros_try_catch
    def read_position(self):
        return self.uarm.get_position()

    @ros_try_catch
    def read_joint_angles(self):
        # return [self.get_joint_angles(), time.time()]
        return self.get_joint_angles()

    def uarm_read(self):
        rospy.loginfo("uarm_read running")
        while True and (rospy.is_shutdown() is False):
            robot_values = self.get_from_read_queue(blocking=False)
            if robot_values == "SHUTDOWN":
                rospy.logwarn("sending shutdown signal")
                rospy.signal_shutdown("Normal Shutdown Procedure")
                break
            elif not robot_values:
                pass
            elif robot_values == "DONE":
                self.uarm_read_pub.publish("DONE")
            else:
                msg = str(robot_values)
                msg = msg.translate(None, '[]')
                self.uarm_read_pub.publish(msg)
                #self.ros_rate.sleep()

        rospy.loginfo("uarm_read shutdown")

    def send_to_interface_queue(self, msg, msg_list=False, priority=10):
        with self.interface_queue_put_lock:
            if msg_list:
                for m in msg_list:
                    self.uarm_interface_queue.put((priority,m))
            else:
                self.uarm_interface_queue.put((priority, msg))

    # def send_to_interface_queue(self, msg, msg_list=False, priority=False):
    #     with self.interface_queue_put_lock:
    #         if msg_list:
    #             if priority:
    #                 msg_list = msg + self.get_from_interface_queue(all_msgs=True)
    #             for m in msg_list:
    #                 self.uarm_interface_queue.put(m)
    #         else:
    #             if priority:
    #                 msg_list = [msg] + self.get_from_interface_queue(all_msgs=True)
    #                 for m in msg_list:
    #                     self.uarm_interface_queue.put(m)
    #             else:
    #                 self.uarm_interface_queue.put(msg)

    def filter_interface_queue(self, msg_filter, invert=False):
        msgs = self.get_from_interface_queue(all_msgs=True, blocking=True)
        if msg_filter is not "":
            # print "FILTER: ", filter
            n = 0
            for i, msg in enumerate(msgs):
                if filter in msg and not invert:
                    msgs.pop(i)
                    n += 1
                elif filter not in msg and invert:
                    msgs.pop(i)
                    n += 1
            if msgs is not []:
                # print n, " messages filtered matching search terms", filter, invert
                # print msgs
                #self.send_to_interface_queue(msgs, msg_list=True, priority=True)
                self.send_to_interface_queue(msgs, msg_list=True, priority=1)

    def get_from_interface_queue(self, all_msgs=False, blocking=False):
        with self.interface_queue_get_lock:
            if all_msgs:
                msgs = []
                while self.uarm_interface_queue.qsize() > 0:
                    try:
                        msg = self.uarm_interface_queue.get(blocking)
                        try:
                            msg = msg[1]
                        except:
                            pass
                        print msg
                        msgs.append(msg)

                    except Queue.Empty:
                        msgs = []

                return msgs
            else:
                try:
                    msg = self.uarm_interface_queue.get(blocking)
                except Queue.Empty:
                    msg = ""
                try:
                    msg = msg[1]
                except:
                    pass
                return msg

    def send_to_read_queue(self, msg, msg_list=False, priority=10):
        with self.read_queue_put_lock:
            if msg_list:
                for m in msg:
                    self.uarm_read_queue.put((priority,m))
            else:
                self.uarm_read_queue.put((priority,msg))

    # def send_to_read_queue(self, msg, msg_list=False, priority=False):
    #     with self.read_queue_put_lock:
    #         if msg_list:
    #             if priority:
    #                 msg_list = msg + self.get_from_read_queue(all_msgs=True)
    #             for m in msg_list:
    #                 self.uarm_read_queue.put(m)
    #         else:
    #             if priority:
    #                 # print "Sending priority message, read_queue:", msg
    #                 msg_list = [msg] + self.get_from_read_queue(all_msgs=True)
    #                 for m in msg_list:
    #                     self.uarm_read_queue.put(m)
    #             else:
    #                 self.uarm_read_queue.put(msg)

    def get_from_read_queue(self, all_msgs=False, blocking=False):
        with self.read_queue_get_lock:
            if all_msgs:
                msgs = []
                while self.uarm_read_queue.qsize() > 0:
                    try:
                        msg = self.uarm_read_queue.get(blocking)
                        try:
                            msg = msg[1:]
                        except:
                            pass
                        msgs.append()
                    except Queue.Empty:
                        msgs = []
                return msgs
            else:
                try:
                    msg = self.uarm_read_queue.get(blocking)
                    try:
                        msg = msg[1:]
                    except:
                        pass

                except Queue.Empty:
                    msg = ""
                return msg

    def uarm_write_callback(self, data):
        # print "DATA: ", data.data
        if data.data == "CLEAR":
                self.get_from_interface_queue(all_msgs=True)
                self.get_from_read_queue(all_msgs=True)

        elif data.data[0] == "!":
                msg = data.data[1:]
                #print msg
                #self.send_to_interface_queue(msg, priority=True)
                self.send_to_interface_queue(msg, priority=1)
        else:
                self.send_to_interface_queue(data.data)

    def playback(self):
        if self.playback_data:
            self.loading = False
            self.get_from_interface_queue(all_msgs=True)  # flush queue
            try:
                joint_angles = self.ja
                self.send_to_interface_queue("JA"+str(joint_angles[0]) + "," +
                                             str(joint_angles[1]) + "," +
                                             str(joint_angles[2]) +
                                             str(joint_angles[3]))
            except Exception as e:
                # print "joint angles not ready yet: ", e
                err = str(e.message)
                rospy.logerr("joint angles not ready yet: " + err)
                rospy.signal_shutdown("joint angles not ready yet")

            self.send_to_interface_queue("ATT")

            self.filter_interface_queue("READ")

            for line in self.playback_data:
                # print self.playback_active, line
                if self.playback_active is False:
                    self.get_from_interface_queue(all_msgs=True)
                    break
                msg = str(line)

                self.send_to_interface_queue("READ")
                self.send_to_interface_queue(msg)

        self.send_to_interface_queue("BEEP")
        self.send_to_interface_queue("DET")
        self.playback_data = []

    def load_playback_data(self, data_line):
        self.loading = True
        rospy.loginfo("Loading playback data...")
        self.playback_data.append(data_line)

    @ros_try_catch
    def process_command(self, command):
        if command == "PUMP_ON":
            self.uarm.set_pump(1)

        if command == "PUMP_OFF":
            self.uarm.set_pump(0)

        if command == "ATT":
            self.uarm.set_servo_attach()

        if command == "DET":
            self.uarm.set_servo_detach()

        if command[0:3] == "POS" and self.playback_active:
            position = map(float, command[3:].split(','))
            self.uarm.set_position(position[0], position[1], position[2])

        if command[0:3] == "REL":
            adjust = float(command[3:])
            # print adjust
            self.uarm.set_position(x=0, y=0, z=adjust, speed=0, relative=True)

        if command[0:2] == "JA" and self.playback_active:
            angle = map(float, command[2:].split(','))
            for i in range(0, len(angle)-1):
                self.uarm.set_servo_angle(i, angle[i])

        if command[0:4] == "LOAD":
            self.load_playback_data(command[4:])

        if command == "PLAY":
            if self.playback_data:
                self.playback_active = True
                self.playback_thread = threading.Thread(target=self.playback)
                self.playback_thread.daemon = True

                self.playback_thread.start()
        if command == "BEEP":
            self.uarm.set_buzzer(10000, 0.1)
            #self.send_to_read_queue("DONE", priority=True)
            self.send_to_read_queue("DONE", priority=1)
            self.playback_active = False

        if command == "STOP":
            self.playback_active = False
            self.filter_interface_queue(msg_filter="READ")
            time.sleep(0.5)
            self.filter_interface_queue(msg_filter="JA")
            self.playback_data = []
<|endoftext|>


<commit_before>""" support for providing temporary directories to test functions.  """
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import re
import tempfile
import warnings

import attr
import py
import six

import pytest
from .pathlib import ensure_reset_dir
from .pathlib import LOCK_TIMEOUT
from .pathlib import make_numbered_dir
from .pathlib import make_numbered_dir_with_cleanup
from .pathlib import Path
from _pytest.monkeypatch import MonkeyPatch

@attr.s
class TempPathFactory(object):
    """Factory for temporary directories under the common base temp directory.

    The base directory can be configured using the ``--basetemp`` option."""

    _given_basetemp = attr.ib(
        # using os.path.abspath() to get absolute path instead of resolve() as it
        # does not work the same in all platforms (see #4427)
        # Path.absolute() exists, but it is not public (see https://bugs.python.org/issue25012)
        convert=attr.converters.optional(
            lambda p: Path(os.path.abspath(six.text_type(p)))
        )
    )
    _trace = attr.ib()
    _basetemp = attr.ib(default=None)

    @classmethod
    def from_config(cls, config):
        """
        :param config: a pytest configuration
        """
        return cls(
            given_basetemp=config.option.basetemp, trace=config.trace.get("tmpdir")
        )

    def mktemp(self, basename, numbered=True):
        """makes a temporary directory managed by the factory"""
        if not numbered:
            p = self.getbasetemp().joinpath(basename)
            p.mkdir()
        else:
            p = make_numbered_dir(root=self.getbasetemp(), prefix=basename)
            self._trace("mktemp", p)
        return p

    def getbasetemp(self):
        """ return base temporary directory. """
        if self._basetemp is None:
            if self._given_basetemp is not None:
                basetemp = self._given_basetemp
                ensure_reset_dir(basetemp)
                basetemp = basetemp.resolve()
            else:
                from_env = os.environ.get("PYTEST_DEBUG_TEMPROOT")
                temproot = Path(from_env or tempfile.gettempdir()).resolve()
                user = get_user() or "unknown"
                # use a sub-directory in the temproot to speed-up
                # make_numbered_dir() call
                rootdir = temproot.joinpath("pytest-of-{}".format(user))
                rootdir.mkdir(exist_ok=True)
                basetemp = make_numbered_dir_with_cleanup(
                    prefix="pytest-", root=rootdir, keep=3, lock_timeout=LOCK_TIMEOUT
                )
            assert basetemp is not None
            self._basetemp = t = basetemp
            self._trace("new basetemp", t)
            return t
        else:
            return self._basetemp

@attr.s
class TempdirFactory(object):
    """
    backward comptibility wrapper that implements
    :class:``py.path.local`` for :class:``TempPathFactory``
    """

    _tmppath_factory = attr.ib()

    def ensuretemp(self, string, dir=1):
        """ (deprecated) return temporary directory path with
            the given string as the trailing part.  It is usually
            better to use the 'tmpdir' function argument which
            provides an empty unique-per-test-invocation directory
            and is guaranteed to be empty.
        """
        # py.log._apiwarn(">1.1", "use tmpdir function argument")
        from .deprecated import PYTEST_ENSURETEMP

        warnings.warn(PYTEST_ENSURETEMP, stacklevel=2)
        return self.getbasetemp().ensure(string, dir=dir)

    def mktemp(self, basename, numbered=True):
        """Create a subdirectory of the base temporary directory and return it.
        If ``numbered``, ensure the directory is unique by adding a number
        prefix greater than any existing one.
        """
        return py.path.local(self._tmppath_factory.mktemp(basename, numbered).resolve())

    def getbasetemp(self):
        """backward compat wrapper for ``_tmppath_factory.getbasetemp``"""
        return py.path.local(self._tmppath_factory.getbasetemp().resolve())

def get_user():
    """Return the current user name, or None if getuser() does not work
    in the current environment (see #1010).
    """
    import getpass

    try:
        return getpass.getuser()
    except (ImportError, KeyError):
        return None

def pytest_configure(config):
    """Create a TempdirFactory and attach it to the config object.

    This is to comply with existing plugins which expect the handler to be
    available at pytest_configure time, but ideally should be moved entirely
    to the tmpdir_factory session fixture.
    """
    mp = MonkeyPatch()
    tmppath_handler = TempPathFactory.from_config(config)
    t = TempdirFactory(tmppath_handler)
    config._cleanup.append(mp.undo)
    mp.setattr(config, "_tmp_path_factory", tmppath_handler, raising=False)
    mp.setattr(config, "_tmpdirhandler", t, raising=False)
    mp.setattr(pytest, "ensuretemp", t.ensuretemp, raising=False)

@pytest.fixture(scope="session")
def tmpdir_factory(request):
    """Return a :class:`_pytest.tmpdir.TempdirFactory` instance for the test session.
    """
    return request.config._tmpdirhandler

@pytest.fixture(scope="session")
def tmp_path_factory(request):
    """Return a :class:`_pytest.tmpdir.TempPathFactory` instance for the test session.
    """
    return request.config._tmp_path_factory

def _mk_tmp(request, factory):
    name = request.node.name
    name = re.sub(r"[\W]", "_", name)
    MAXVAL = 30
    name = name[:MAXVAL]
    return factory.mktemp(name, numbered=True)

@pytest.fixture
def tmpdir(tmp_path):
    """Return a temporary directory path object
    which is unique to each test function invocation,
    created as a sub directory of the base temporary
    directory.  The returned object is a `py.path.local`_
    path object.

    .. _`py.path.local`: https://py.readthedocs.io/en/latest/path.html
    """
    return py.path.local(tmp_path)

@pytest.fixture
def tmp_path(request, tmp_path_factory):
    """Return a temporary directory path object
    which is unique to each test function invocation,
    created as a sub directory of the base temporary
    directory.  The returned object is a :class:`pathlib.Path`
    object.

    .. note::

        in python < 3.6 this is a pathlib2.Path
    """

    return _mk_tmp(request, tmp_path_factory)
<commit_msg>Update kwarg for attr.ib to use 'converter' as 'convert' is due to be deprecated<commit_after>""" support for providing temporary directories to test functions.  """
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import re
import tempfile
import warnings

import attr
import py
import six

import pytest
from .pathlib import ensure_reset_dir
from .pathlib import LOCK_TIMEOUT
from .pathlib import make_numbered_dir
from .pathlib import make_numbered_dir_with_cleanup
from .pathlib import Path
from _pytest.monkeypatch import MonkeyPatch

@attr.s
class TempPathFactory(object):
    """Factory for temporary directories under the common base temp directory.

    The base directory can be configured using the ``--basetemp`` option."""

    _given_basetemp = attr.ib(
        # using os.path.abspath() to get absolute path instead of resolve() as it
        # does not work the same in all platforms (see #4427)
        # Path.absolute() exists, but it is not public (see https://bugs.python.org/issue25012)
        converter=attr.converters.optional(
            lambda p: Path(os.path.abspath(six.text_type(p)))
        )
    )
    _trace = attr.ib()
    _basetemp = attr.ib(default=None)

    @classmethod
    def from_config(cls, config):
        """
        :param config: a pytest configuration
        """
        return cls(
            given_basetemp=config.option.basetemp, trace=config.trace.get("tmpdir")
        )

    def mktemp(self, basename, numbered=True):
        """makes a temporary directory managed by the factory"""
        if not numbered:
            p = self.getbasetemp().joinpath(basename)
            p.mkdir()
        else:
            p = make_numbered_dir(root=self.getbasetemp(), prefix=basename)
            self._trace("mktemp", p)
        return p

    def getbasetemp(self):
        """ return base temporary directory. """
        if self._basetemp is None:
            if self._given_basetemp is not None:
                basetemp = self._given_basetemp
                ensure_reset_dir(basetemp)
                basetemp = basetemp.resolve()
            else:
                from_env = os.environ.get("PYTEST_DEBUG_TEMPROOT")
                temproot = Path(from_env or tempfile.gettempdir()).resolve()
                user = get_user() or "unknown"
                # use a sub-directory in the temproot to speed-up
                # make_numbered_dir() call
                rootdir = temproot.joinpath("pytest-of-{}".format(user))
                rootdir.mkdir(exist_ok=True)
                basetemp = make_numbered_dir_with_cleanup(
                    prefix="pytest-", root=rootdir, keep=3, lock_timeout=LOCK_TIMEOUT
                )
            assert basetemp is not None
            self._basetemp = t = basetemp
            self._trace("new basetemp", t)
            return t
        else:
            return self._basetemp

@attr.s
class TempdirFactory(object):
    """
    backward comptibility wrapper that implements
    :class:``py.path.local`` for :class:``TempPathFactory``
    """

    _tmppath_factory = attr.ib()

    def ensuretemp(self, string, dir=1):
        """ (deprecated) return temporary directory path with
            the given string as the trailing part.  It is usually
            better to use the 'tmpdir' function argument which
            provides an empty unique-per-test-invocation directory
            and is guaranteed to be empty.
        """
        # py.log._apiwarn(">1.1", "use tmpdir function argument")
        from .deprecated import PYTEST_ENSURETEMP

        warnings.warn(PYTEST_ENSURETEMP, stacklevel=2)
        return self.getbasetemp().ensure(string, dir=dir)

    def mktemp(self, basename, numbered=True):
        """Create a subdirectory of the base temporary directory and return it.
        If ``numbered``, ensure the directory is unique by adding a number
        prefix greater than any existing one.
        """
        return py.path.local(self._tmppath_factory.mktemp(basename, numbered).resolve())

    def getbasetemp(self):
        """backward compat wrapper for ``_tmppath_factory.getbasetemp``"""
        return py.path.local(self._tmppath_factory.getbasetemp().resolve())

def get_user():
    """Return the current user name, or None if getuser() does not work
    in the current environment (see #1010).
    """
    import getpass

    try:
        return getpass.getuser()
    except (ImportError, KeyError):
        return None

def pytest_configure(config):
    """Create a TempdirFactory and attach it to the config object.

    This is to comply with existing plugins which expect the handler to be
    available at pytest_configure time, but ideally should be moved entirely
    to the tmpdir_factory session fixture.
    """
    mp = MonkeyPatch()
    tmppath_handler = TempPathFactory.from_config(config)
    t = TempdirFactory(tmppath_handler)
    config._cleanup.append(mp.undo)
    mp.setattr(config, "_tmp_path_factory", tmppath_handler, raising=False)
    mp.setattr(config, "_tmpdirhandler", t, raising=False)
    mp.setattr(pytest, "ensuretemp", t.ensuretemp, raising=False)

@pytest.fixture(scope="session")
def tmpdir_factory(request):
    """Return a :class:`_pytest.tmpdir.TempdirFactory` instance for the test session.
    """
    return request.config._tmpdirhandler

@pytest.fixture(scope="session")
def tmp_path_factory(request):
    """Return a :class:`_pytest.tmpdir.TempPathFactory` instance for the test session.
    """
    return request.config._tmp_path_factory

def _mk_tmp(request, factory):
    name = request.node.name
    name = re.sub(r"[\W]", "_", name)
    MAXVAL = 30
    name = name[:MAXVAL]
    return factory.mktemp(name, numbered=True)

@pytest.fixture
def tmpdir(tmp_path):
    """Return a temporary directory path object
    which is unique to each test function invocation,
    created as a sub directory of the base temporary
    directory.  The returned object is a `py.path.local`_
    path object.

    .. _`py.path.local`: https://py.readthedocs.io/en/latest/path.html
    """
    return py.path.local(tmp_path)

@pytest.fixture
def tmp_path(request, tmp_path_factory):
    """Return a temporary directory path object
    which is unique to each test function invocation,
    created as a sub directory of the base temporary
    directory.  The returned object is a :class:`pathlib.Path`
    object.

    .. note::

        in python < 3.6 this is a pathlib2.Path
    """

    return _mk_tmp(request, tmp_path_factory)
<|endoftext|>


<commit_before>from elasticsearch_dsl.connections import connections
from elasticsearch import helpers
from elasticsearch_dsl import Search, Index
from elasticsearch.exceptions import ConnectionError
from glob import glob
import gzip
import logging
import json
import argparse
from datetime import datetime
from time import sleep
import os
import itertools
import math

from models import TweetIndex, to_tweet, DatasetIndex, to_dataset, DatasetDocType, get_tweets_index_name
from utils import read_json, short_uid

log = logging.getLogger(__name__)

connections.create_connection(hosts=['elasticsearch'], timeout=90, sniff_on_start=True, sniff_on_connection_fail=True,
                              retry_on_timeout=True)

CONNECTION_ERROR_TRIES = 30

def find_files(path):
    """
    Returns (.json files, .json.gz files, .txt files) found in path.
    """
    json_filepaths = glob('{}/*.json'.format(path))
    dataset_filepath = os.path.join(path, 'dataset.json')
    if dataset_filepath in json_filepaths:
        json_filepaths.remove(dataset_filepath)
    return (json_filepaths,
            glob('{}/*.json.gz'.format(path)),
            glob('{}/*.txt'.format(path)))

def count_lines(json_files, json_gz_files, txt_files):
    total_lines = 0
    for filepath in json_files:
        total_lines += sum(1 for _ in open(filepath))
    for filepath in json_gz_files:
        total_lines += sum(1 for _ in gzip.open(filepath))
    for filepath in txt_files:
        total_lines += sum(1 for _ in open(filepath))
    return total_lines

def count_files(json_files, json_gz_files, txt_files):
    return len(json_files) + len(json_gz_files) + len(txt_files)

def tweet_iter(json_files, json_gz_files, txt_files, limit=None, total_tweets=None):
    counter = 0
    for filepath in json_files:
        with open(filepath) as file:
            for line in file:
                if counter % 10000 == 0:
                    log.info('{:,} of {:,} tweets'.format(counter, limit or total_tweets or 0))
                if counter == limit:
                    break
                counter += 1
                yield json.loads(line)
    for filepath in json_gz_files:
        with gzip.open(filepath) as file:
            for line in file:
                if counter % 10000 == 0:
                    log.info('{:,} of {:,} tweets'.format(counter, limit or total_tweets or 0))
                if counter == limit:
                    break
                counter += 1
                yield json.loads(line)

                # TODO: Handle hydration

def delete_tweet_index(dataset_identifier):
    tweet_index = TweetIndex(dataset_identifier)
    tweet_index.delete(ignore=404)
    log.info('Deleted tweets from {}'.format(dataset_identifier))

def get_tweet_index_state(dataset_identifier):
    """
    Returns alias name, new index name, existing index name or none
    """
    alias_name = get_tweets_index_name(dataset_identifier)
    index_name_a = '{}-a'.format(alias_name)
    index_name_b = '{}-b'.format(alias_name)
    # If a exists then a is existing and b is new
    if Index(index_name_a).exists():
        return alias_name, index_name_b, index_name_a
    # Else if b exists then a is new and b is existing
    elif Index(index_name_b).exists():
        return alias_name, index_name_b, index_name_a
    # Else neither exists and a is new
    return alias_name, index_name_a, None

def update_dataset_stats(dataset):
    search = Search(index=get_tweets_index_name(dataset.meta.id))
    search = search.query('term', dataset_id=dataset.meta.id)[0:0]
    search.aggs.metric('created_at_min', 'min', field='created_at')
    search.aggs.metric('created_at_max', 'max', field='created_at')
    search_response = search.execute()
    dataset.first_tweet_created_at = datetime.utcfromtimestamp(
        search_response.aggregations.created_at_min.value / 1000.0)
    dataset.last_tweet_created_at = datetime.utcfromtimestamp(
        search_response.aggregations.created_at_max.value / 1000.0)
    dataset.tweet_count = search_response.hits.total
    dataset.save()

if __name__ == '__main__':
    parser = argparse.ArgumentParser('tweetset_loader')
    parser.add_argument('--debug', action='store_true')

    # Subparsers
    subparsers = parser.add_subparsers(dest='command', help='command help')

    update_parser = subparsers.add_parser('update', help='update dataset metadata')
    update_parser.add_argument('dataset_identifier', help='identifier (a UUID) for the dataset')
    update_parser.add_argument('--path', help='path of dataset', default='/dataset')
    update_parser.add_argument('--filename', help='filename of dataset file', default='dataset.json')
    update_parser.add_argument('--stats', action='store_true', help='Also update dataset statistics')
    update_parser.add_argument('--create', action='store_true', help='Create if does not exist.')

    delete_parser = subparsers.add_parser('delete', help='delete dataset and tweets')
    delete_parser.add_argument('dataset_identifier', help='identifier (a UUID) for the dataset')

    reload_parser = subparsers.add_parser('reload', help='reload the tweets in a dataset')
    reload_parser.add_argument('dataset_identifier', help='identifier (a UUID) for the dataset')
    reload_parser.add_argument('--path', help='path of the directory containing the tweet files', default='/dataset')
    reload_parser.add_argument('--limit', type=int, help='limit the number of tweets to load')
    reload_parser.add_argument('--skip-count', action='store_true', help='skip count the tweets')
    reload_parser.add_argument('--store-tweet', action='store_true', help='store the entire tweet')
    reload_parser.add_argument('--replicas', type=int, default='1', help='number of replicas to make of this dataset')

    dataset_parser = subparsers.add_parser('create', help='create a dataset and add tweets')
    dataset_parser.add_argument('--path', help='path of dataset', default='/dataset')
    dataset_parser.add_argument('--filename', help='filename of dataset file', default='dataset.json')
    dataset_parser.add_argument('--limit', type=int, help='limit the number of tweets to load')
    dataset_parser.add_argument('--skip-count', action='store_true', help='skip count the tweets')
    dataset_parser.add_argument('--store-tweet', action='store_true', help='store the entire tweet')
    dataset_parser.add_argument('--shards', type=int, help='number of shards for this dataset')
    dataset_parser.add_argument('--replicas', type=int, default='1', help='number of replicas to make of this dataset')

    subparsers.add_parser('clear', help='delete all indexes')

    args = parser.parse_args()

    # Logging
    logging.getLogger('urllib3').setLevel(logging.WARNING)
    logging.getLogger('elasticsearch').setLevel(logging.WARNING)
    logging.basicConfig(
        level=logging.DEBUG if args.debug else logging.INFO
    )

    if not args.command:
        parser.print_help()
        quit()

    # Create indexs if they doesn't exist
    dataset_index = DatasetIndex()
    dataset_index.create(ignore=400)

    dataset_id = None
    if args.command == 'delete':
        dataset = DatasetDocType.get(args.dataset_identifier)
        if not dataset:
            raise Exception('{} not found'.format(args.dataset_identifier))
        dataset.delete()
        log.info('Deleted {}'.format(dataset.meta.id))
        delete_tweet_index(args.dataset_identifier)
    if args.command == 'create':
        dataset = to_dataset(read_json(os.path.join(args.path, args.filename)),
                             dataset_id=short_uid(6,
                                                  exists_func=lambda uid: DatasetDocType.get(uid,
                                                                                ignore=404) is not None))
        dataset.save()
        dataset_id = dataset.meta.id
        log.info('Created {}'.format(dataset_id))
        print('Dataset id is {}'.format(dataset_id))

    if args.command in ('create', 'reload'):
        if dataset_id is None:
            dataset_id = args.dataset_identifier
        store_tweet = os.environ.get('STORE_TWEET', 'false').lower() == 'true' or args.store_tweet
        if store_tweet:
            log.info('Storing tweet')
        dataset = DatasetDocType.get(dataset_id)
        if not dataset:
            raise Exception('{} not found'.format(dataset_id))

        filepaths = find_files(args.path)
        file_count = count_files(*filepaths)
        tweet_count = 0
        if not args.skip_count:
            log.info('Counting tweets in %s files.', file_count)
            tweet_count = count_lines(*filepaths)
            log.info('{:,} total tweets'.format(tweet_count))

        # Create the index
        # In testing, 500K tweets (storing tweet) = 615MB
        # Thus, 32.5 million tweets per shard to have a max shard size of 40GB
        # In testing, 500k tweets (not storing tweet) = 145MB
        # Thus, 138 million tweets per shard to have a max shard size of 40GB
        tweets_per_shard = 32500000 if store_tweet else 138000000
        shards = (args.shards if hasattr(args, 'shards') else None) or math.ceil(
            float(tweet_count) / tweets_per_shard) or 1
        log.info('Using %s shards and %s replicas for index.', shards, args.replicas)

        alias_name, new_index_name, existing_index_name = get_tweet_index_state(dataset_id)

        log.debug('Index name is %s', new_index_name)
        tweet_index = TweetIndex(new_index_name, shards=shards, replicas=0, refresh_interval=-1)
        tweet_index.create()

        # create_tweet_index(dataset_id, shards, replicas=0, )
        connection = connections.get_connection()

        helpers.bulk(connection,
                     [to_tweet(tweet_json, dataset_id, new_index_name, store_tweet=store_tweet).to_dict(
                         include_meta=True) for
                      tweet_json in
                      tweet_iter(*filepaths, limit=args.limit, total_tweets=tweet_count)],
                     max_retries=CONNECTION_ERROR_TRIES)

        log.debug('Setting replicas and refresh interval')
        tweet_index.put_settings(body={
            'number_of_replicas': args.replicas, 'refresh_interval': '1s'})

        # Add aliases
        log.debug('Adding alias %s to %s', alias_name, new_index_name)
        tweet_index.put_alias(name=alias_name)

        # Delete existing index
        if existing_index_name:
            log.debug('Deleting existing index %s', existing_index_name)
            TweetIndex(existing_index_name).delete()

        # Get number of tweets in dataset and update
        sleep(10)
        update_dataset_stats(dataset)

    if args.command == 'clear':
        search = DatasetDocType.search()
        for dataset in search.execute():
            delete_tweet_index(dataset.meta.id)
        dataset_index.delete()
        log.info("Deleted indexes")
<commit_msg>Fixed deleting indexes.<commit_after>from elasticsearch_dsl.connections import connections
from elasticsearch import helpers
from elasticsearch_dsl import Search, Index
from elasticsearch.exceptions import ConnectionError
from glob import glob
import gzip
import logging
import json
import argparse
from datetime import datetime
from time import sleep
import os
import itertools
import math

from models import TweetIndex, to_tweet, DatasetIndex, to_dataset, DatasetDocType, get_tweets_index_name
from utils import read_json, short_uid

log = logging.getLogger(__name__)

connections.create_connection(hosts=['elasticsearch'], timeout=90, sniff_on_start=True, sniff_on_connection_fail=True,
                              retry_on_timeout=True)

CONNECTION_ERROR_TRIES = 30

def find_files(path):
    """
    Returns (.json files, .json.gz files, .txt files) found in path.
    """
    json_filepaths = glob('{}/*.json'.format(path))
    dataset_filepath = os.path.join(path, 'dataset.json')
    if dataset_filepath in json_filepaths:
        json_filepaths.remove(dataset_filepath)
    return (json_filepaths,
            glob('{}/*.json.gz'.format(path)),
            glob('{}/*.txt'.format(path)))

def count_lines(json_files, json_gz_files, txt_files):
    total_lines = 0
    for filepath in json_files:
        total_lines += sum(1 for _ in open(filepath))
    for filepath in json_gz_files:
        total_lines += sum(1 for _ in gzip.open(filepath))
    for filepath in txt_files:
        total_lines += sum(1 for _ in open(filepath))
    return total_lines

def count_files(json_files, json_gz_files, txt_files):
    return len(json_files) + len(json_gz_files) + len(txt_files)

def tweet_iter(json_files, json_gz_files, txt_files, limit=None, total_tweets=None):
    counter = 0
    for filepath in json_files:
        with open(filepath) as file:
            for line in file:
                if counter % 10000 == 0:
                    log.info('{:,} of {:,} tweets'.format(counter, limit or total_tweets or 0))
                if counter == limit:
                    break
                counter += 1
                yield json.loads(line)
    for filepath in json_gz_files:
        with gzip.open(filepath) as file:
            for line in file:
                if counter % 10000 == 0:
                    log.info('{:,} of {:,} tweets'.format(counter, limit or total_tweets or 0))
                if counter == limit:
                    break
                counter += 1
                yield json.loads(line)

                # TODO: Handle hydration

def delete_tweet_index(dataset_identifier):
    _, __, existing_index_name = get_tweet_index_state(dataset_identifier)
    if existing_index_name:
        TweetIndex(existing_index_name).delete(ignore=404)
        log.info('Deleted tweets from {}'.format(dataset_identifier))

def get_tweet_index_state(dataset_identifier):
    """
    Returns alias name, new index name, existing index name or none
    """
    alias_name = get_tweets_index_name(dataset_identifier)
    index_name_a = '{}-a'.format(alias_name)
    index_name_b = '{}-b'.format(alias_name)
    # If a exists then a is existing and b is new
    if Index(index_name_a).exists():
        return alias_name, index_name_b, index_name_a
    # Else if b exists then a is new and b is existing
    elif Index(index_name_b).exists():
        return alias_name, index_name_b, index_name_a
    # Else neither exists and a is new
    return alias_name, index_name_a, None

def update_dataset_stats(dataset):
    search = Search(index=get_tweets_index_name(dataset.meta.id))
    search = search.query('term', dataset_id=dataset.meta.id)[0:0]
    search.aggs.metric('created_at_min', 'min', field='created_at')
    search.aggs.metric('created_at_max', 'max', field='created_at')
    search_response = search.execute()
    dataset.first_tweet_created_at = datetime.utcfromtimestamp(
        search_response.aggregations.created_at_min.value / 1000.0)
    dataset.last_tweet_created_at = datetime.utcfromtimestamp(
        search_response.aggregations.created_at_max.value / 1000.0)
    dataset.tweet_count = search_response.hits.total
    dataset.save()

if __name__ == '__main__':
    parser = argparse.ArgumentParser('tweetset_loader')
    parser.add_argument('--debug', action='store_true')

    # Subparsers
    subparsers = parser.add_subparsers(dest='command', help='command help')

    update_parser = subparsers.add_parser('update', help='update dataset metadata')
    update_parser.add_argument('dataset_identifier', help='identifier (a UUID) for the dataset')
    update_parser.add_argument('--path', help='path of dataset', default='/dataset')
    update_parser.add_argument('--filename', help='filename of dataset file', default='dataset.json')
    update_parser.add_argument('--stats', action='store_true', help='Also update dataset statistics')
    update_parser.add_argument('--create', action='store_true', help='Create if does not exist.')

    delete_parser = subparsers.add_parser('delete', help='delete dataset and tweets')
    delete_parser.add_argument('dataset_identifier', help='identifier (a UUID) for the dataset')

    reload_parser = subparsers.add_parser('reload', help='reload the tweets in a dataset')
    reload_parser.add_argument('dataset_identifier', help='identifier (a UUID) for the dataset')
    reload_parser.add_argument('--path', help='path of the directory containing the tweet files', default='/dataset')
    reload_parser.add_argument('--limit', type=int, help='limit the number of tweets to load')
    reload_parser.add_argument('--skip-count', action='store_true', help='skip count the tweets')
    reload_parser.add_argument('--store-tweet', action='store_true', help='store the entire tweet')
    reload_parser.add_argument('--replicas', type=int, default='1', help='number of replicas to make of this dataset')

    dataset_parser = subparsers.add_parser('create', help='create a dataset and add tweets')
    dataset_parser.add_argument('--path', help='path of dataset', default='/dataset')
    dataset_parser.add_argument('--filename', help='filename of dataset file', default='dataset.json')
    dataset_parser.add_argument('--limit', type=int, help='limit the number of tweets to load')
    dataset_parser.add_argument('--skip-count', action='store_true', help='skip count the tweets')
    dataset_parser.add_argument('--store-tweet', action='store_true', help='store the entire tweet')
    dataset_parser.add_argument('--shards', type=int, help='number of shards for this dataset')
    dataset_parser.add_argument('--replicas', type=int, default='1', help='number of replicas to make of this dataset')

    subparsers.add_parser('clear', help='delete all indexes')

    args = parser.parse_args()

    # Logging
    logging.getLogger('urllib3').setLevel(logging.WARNING)
    logging.getLogger('elasticsearch').setLevel(logging.WARNING)
    logging.basicConfig(
        level=logging.DEBUG if args.debug else logging.INFO
    )

    if not args.command:
        parser.print_help()
        quit()

    # Create indexs if they doesn't exist
    dataset_index = DatasetIndex()
    dataset_index.create(ignore=400)

    dataset_id = None
    if args.command == 'delete':
        dataset = DatasetDocType.get(args.dataset_identifier)
        if not dataset:
            raise Exception('{} not found'.format(args.dataset_identifier))
        dataset.delete()
        log.info('Deleted {}'.format(dataset.meta.id))
        delete_tweet_index(args.dataset_identifier)
    if args.command == 'create':
        dataset = to_dataset(read_json(os.path.join(args.path, args.filename)),
                             dataset_id=short_uid(6,
                                                  exists_func=lambda uid: DatasetDocType.get(uid,
                                                                                ignore=404) is not None))
        dataset.save()
        dataset_id = dataset.meta.id
        log.info('Created {}'.format(dataset_id))
        print('Dataset id is {}'.format(dataset_id))

    if args.command in ('create', 'reload'):
        if dataset_id is None:
            dataset_id = args.dataset_identifier
        store_tweet = os.environ.get('STORE_TWEET', 'false').lower() == 'true' or args.store_tweet
        if store_tweet:
            log.info('Storing tweet')
        dataset = DatasetDocType.get(dataset_id)
        if not dataset:
            raise Exception('{} not found'.format(dataset_id))

        filepaths = find_files(args.path)
        file_count = count_files(*filepaths)
        tweet_count = 0
        if not args.skip_count:
            log.info('Counting tweets in %s files.', file_count)
            tweet_count = count_lines(*filepaths)
            log.info('{:,} total tweets'.format(tweet_count))

        # Create the index
        # In testing, 500K tweets (storing tweet) = 615MB
        # Thus, 32.5 million tweets per shard to have a max shard size of 40GB
        # In testing, 500k tweets (not storing tweet) = 145MB
        # Thus, 138 million tweets per shard to have a max shard size of 40GB
        tweets_per_shard = 32500000 if store_tweet else 138000000
        shards = (args.shards if hasattr(args, 'shards') else None) or math.ceil(
            float(tweet_count) / tweets_per_shard) or 1
        log.info('Using %s shards and %s replicas for index.', shards, args.replicas)

        alias_name, new_index_name, existing_index_name = get_tweet_index_state(dataset_id)

        log.debug('Index name is %s', new_index_name)
        tweet_index = TweetIndex(new_index_name, shards=shards, replicas=0, refresh_interval=-1)
        tweet_index.create()

        # create_tweet_index(dataset_id, shards, replicas=0, )
        connection = connections.get_connection()

        helpers.bulk(connection,
                     [to_tweet(tweet_json, dataset_id, new_index_name, store_tweet=store_tweet).to_dict(
                         include_meta=True) for
                      tweet_json in
                      tweet_iter(*filepaths, limit=args.limit, total_tweets=tweet_count)],
                     max_retries=CONNECTION_ERROR_TRIES)

        log.debug('Setting replicas and refresh interval')
        tweet_index.put_settings(body={
            'number_of_replicas': args.replicas, 'refresh_interval': '1s'})

        # Add aliases
        log.debug('Adding alias %s to %s', alias_name, new_index_name)
        tweet_index.put_alias(name=alias_name)

        # Delete existing index
        if existing_index_name:
            log.debug('Deleting existing index %s', existing_index_name)
            TweetIndex(existing_index_name).delete()

        # Get number of tweets in dataset and update
        sleep(10)
        update_dataset_stats(dataset)

    if args.command == 'clear':
        search = DatasetDocType.search()
        for dataset in search.execute():
            delete_tweet_index(dataset.meta.id)
        dataset_index.delete()
        log.info("Deleted indexes")
<|endoftext|>


<commit_before>import time
from datetime import timedelta

###
### Time Management
###

# Use the function currentTime in this module to get a number representing the current time. The time
# will be denoted in seconds but the return value is not suitable for wall clock time measurement
# but only for measuring time intervals.
#
# For Python 3.3 and later, we will use the CPU performance counter (see PEP-418) which is a monotonic
# clock. For older Python versions we fall back to time.time() which is less accurate and also not
# monotonic.
try:
    currentTime = time.perf_counter
except (AttributeError, NameError):
    currentTime = time.time

def timePeriodSeconds(basis):
    if isinstance(basis, timedelta):
        try:
            return basis.total_seconds()
        except AttributeError:
            # Must be Python 2.6... which doesn't have total_seconds yet
            return (basis.days * 24.0 * 60 * 60) + basis.seconds + (basis.microseconds / 1000.0 / 1000)
    elif isinstance(basis, (int, float)):
        # assume the raw value is already in seconds
        return basis
    raise TypeError('Cannot determine time from a %s argument'%str(type(basis)))

def toTimeDeltaOrNone(timespec):
    if timespec is None: return None
    if isinstance(timespec, timedelta): return timespec
    if isinstance(timespec, int): return timedelta(seconds=timespec)
    if isinstance(timespec, float):
        return timedelta(seconds=int(timespec),
                         microseconds = int((timespec - int(timespec)) * 1000 * 1000))
    raise TypeError('Unknown type for timespec: %s'%type(timespec))

class Timer(object):
    """
    Keeps track of the elapsed time since its creation or the last call to reset().
    """
    def __init__(self):
        self._start = currentTime()
    def elapsed(self):
        """
        :return: a timedelta object representing the elapsed time.
        """
        return timedelta(seconds=(currentTime() - self._start))
    def elapsedSeconds(self):
        """
        :return: A float representing the elapsed time in seconds.
        """
        return timePeriodSeconds(self.elapsed())
    def reset(self):
        """
        Restarts the timer.
        """
        self._start = currentTime()
    def __str__(self):
        return 'Started_on_' + str(self._start)
    def __eq__(self, o):
        return abs(self._start - o._start) < timedelta(microseconds=1)
    def __lt__(self, o):
        return self._start < o._start
    def __gt__(self, o):
        return not self.__lt__(o)
    def __le__(self, o): return self.__eq__(o) or self.__lt__(o)
    def __ge__(self, o): return self.__eq__(o) or self.__gt__(o)
    def __ne__(self, o): return not self.__eq__(o)
    def __nonzero__(self): return self.elapsedSeconds()

class ExpiryTime(object):
    def __init__(self, duration):
        if duration is None:
            self._time_to_quit = None
        elif isinstance(duration, timedelta):
            self._time_to_quit = currentTime() + timePeriodSeconds(duration)
        else:
            self._time_to_quit = currentTime() + duration
    def expired(self):
        return False if self._time_to_quit is None else (currentTime() >= self._time_to_quit)
    def remaining(self, forever=None):
        return forever if self._time_to_quit is None else \
            (timedelta(seconds=0) if currentTime() > self._time_to_quit else \
             (self._time_to_quit - currentTime()))
    def remainingSeconds(self, forever=None):
        return forever if self._time_to_quit is None else \
            (0 if currentTime() > self._time_to_quit else \
             timePeriodSeconds(self._time_to_quit - currentTime()))
    def __str__(self):
        if self._time_to_quit is None: return 'Forever'
        if self.expired():
            return 'Expired_for_%s'%(currentTime() - self._time_to_quit)
        return 'Expires_in_' + str(self.remaining())
    def __eq__(self, o):
        if isinstance(o, timedelta):
            o = ExpiryTime(o)
        if self._time_to_quit == o._time_to_quit: return True
        if self._time_to_quit == None or o._time_to_quit == None: return False
        if self.expired() and o.expired(): return True
        return abs(self._time_to_quit - o._time_to_quit) < timedelta(microseconds=1)
    def __lt__(self, o):
        try:
            if self._time_to_quit is None and o._time_to_quit is None: return False
        except Exception: pass
        if self._time_to_quit is None: return False
        if isinstance(o, timedelta):
            o = ExpiryTime(o)
        if o._time_to_quit is None: return True
        return self._time_to_quit < o._time_to_quit
    def __gt__(self, o):
        try:
            if self._time_to_quit is None and o._time_to_quit is None: return False
        except Exception: pass
        return not self.__lt__(o)
    def __le__(self, o): return self.__eq__(o) or self.__lt__(o)
    def __ge__(self, o): return self.__eq__(o) or self.__gt__(o)
    def __ne__(self, o): return not self.__eq__(o)
    def __bool__(self): return self.expired()
    def __nonzero__(self): return self.expired()

class ExpirationTimer(object):
    """Keeps track of a duration relative to an original time and
       indicates whether that duration has expired or how much time is
       left before it expires.

       May also be initialized with a duration of None, indicating
       that it should never timeout and that `remaining()` should
       return the forever value (defaulting to None).
    """
    def __init__(self, duration=None):
        self.duration = duration
        if duration is None:
            self._time_to_quit = None
        elif isinstance(duration, timedelta):
            self._time_to_quit = currentTime() + timePeriodSeconds(duration)
        else:
            self._time_to_quit = currentTime() + duration
    def expired(self):
        "Returns true if the indicated duration has passed since this was created."
        return False if self._time_to_quit is None else (currentTime() >= self._time_to_quit)
    def remaining(self, forever=None):
        """Returns a timedelta of remaining time until expiration, or 0 if the
           duration has already expired.  Returns forever if no timeout."""
        return forever if self._time_to_quit is None else \
            timedelta(seconds=self.remainingSeconds())
    def remainingSeconds(self, forever=None):
        """Similar to `remaining()`, but returns an floating point value of the
           number of remaining seconds instead of returning a
           timedelta object.
        """
        return forever if self._time_to_quit is None else \
            max(self._time_to_quit - currentTime(), 0)
    def __str__(self):
        if self._time_to_quit is None: return 'Forever'
        if self.expired():
            return 'Expired_for_%s'%(currentTime() - self._time_to_quit)
        return 'Expires_in_' + str(self.remaining())
    def __eq__(self, o):
        if isinstance(o, timedelta):
            o = ExpiryTime(o)
        if self._time_to_quit == o._time_to_quit: return True
        if self._time_to_quit == None or o._time_to_quit == None: return False
        if self.expired() and o.expired(): return True
        return abs(self._time_to_quit - o._time_to_quit) < timedelta(microseconds=1)
    def __lt__(self, o):
        try:
            if self._time_to_quit is None and o._time_to_quit is None: return False
        except Exception: pass
        if self._time_to_quit is None: return False
        if isinstance(o, timedelta):
            o = ExpiryTime(o)
        if o._time_to_quit is None: return True
        return self._time_to_quit < o._time_to_quit
    def __gt__(self, o):
        try:
            if self._time_to_quit is None and o._time_to_quit is None: return False
        except Exception: pass
        return not self.__lt__(o)
    def __le__(self, o): return self.__eq__(o) or self.__lt__(o)
    def __ge__(self, o): return self.__eq__(o) or self.__gt__(o)
    def __ne__(self, o): return not self.__eq__(o)
    def __bool__(self): return self.expired()
    def __nonzero__(self): return self.expired()
<commit_msg>Return timedelta for remaining() and str() for timers and fix ExpirationTimer __eq__.<commit_after>import time
from datetime import timedelta

###
### Time Management
###

# Use the function currentTime in this module to get a number representing the current time. The time
# will be denoted in seconds but the return value is not suitable for wall clock time measurement
# but only for measuring time intervals.
#
# For Python 3.3 and later, we will use the CPU performance counter (see PEP-418) which is a monotonic
# clock. For older Python versions we fall back to time.time() which is less accurate and also not
# monotonic.
try:
    currentTime = time.perf_counter
except (AttributeError, NameError):
    currentTime = time.time

def timePeriodSeconds(basis):
    if isinstance(basis, timedelta):
        try:
            return basis.total_seconds()
        except AttributeError:
            # Must be Python 2.6... which doesn't have total_seconds yet
            return (basis.days * 24.0 * 60 * 60) + basis.seconds + (basis.microseconds / 1000.0 / 1000)
    elif isinstance(basis, (int, float)):
        # assume the raw value is already in seconds
        return basis
    raise TypeError('Cannot determine time from a %s argument'%str(type(basis)))

def toTimeDeltaOrNone(timespec):
    if timespec is None: return None
    if isinstance(timespec, timedelta): return timespec
    if isinstance(timespec, int): return timedelta(seconds=timespec)
    if isinstance(timespec, float):
        return timedelta(seconds=int(timespec),
                         microseconds = int((timespec - int(timespec)) * 1000 * 1000))
    raise TypeError('Unknown type for timespec: %s'%type(timespec))

class Timer(object):
    """
    Keeps track of the elapsed time since its creation or the last call to reset().
    """
    def __init__(self):
        self._start = currentTime()
    def elapsed(self):
        """
        :return: a timedelta object representing the elapsed time.
        """
        return timedelta(seconds=(currentTime() - self._start))
    def elapsedSeconds(self):
        """
        :return: A float representing the elapsed time in seconds.
        """
        return timePeriodSeconds(self.elapsed())
    def reset(self):
        """
        Restarts the timer.
        """
        self._start = currentTime()
    def __str__(self):
        return 'Started_on_' + str(self._start)
    def __eq__(self, o):
        return abs(self._start - o._start) < timedelta(microseconds=1)
    def __lt__(self, o):
        return self._start < o._start
    def __gt__(self, o):
        return not self.__lt__(o)
    def __le__(self, o): return self.__eq__(o) or self.__lt__(o)
    def __ge__(self, o): return self.__eq__(o) or self.__gt__(o)
    def __ne__(self, o): return not self.__eq__(o)
    def __nonzero__(self): return self.elapsedSeconds()

class ExpiryTime(object):
    def __init__(self, duration):
        if duration is None:
            self._time_to_quit = None
        elif isinstance(duration, timedelta):
            self._time_to_quit = currentTime() + timePeriodSeconds(duration)
        else:
            self._time_to_quit = currentTime() + duration
    def expired(self):
        return False if self._time_to_quit is None else (currentTime() >= self._time_to_quit)
    def remaining(self, forever=None):
        return forever if self._time_to_quit is None else \
            (timedelta(seconds=0) if currentTime() > self._time_to_quit else \
             timedelta(seconds=self._time_to_quit - currentTime()))
    def remainingSeconds(self, forever=None):
        return forever if self._time_to_quit is None else \
            (0 if currentTime() > self._time_to_quit else \
             timePeriodSeconds(self._time_to_quit - currentTime()))
    def __str__(self):
        if self._time_to_quit is None: return 'Forever'
        if self.expired():
            return 'Expired_for_%s' % \
                timedelta(seconds=currentTime() - self._time_to_quit)
        return 'Expires_in_' + str(self.remaining())
    def __eq__(self, o):
        if isinstance(o, timedelta):
            o = ExpiryTime(o)
        if self._time_to_quit == o._time_to_quit: return True
        if self._time_to_quit == None or o._time_to_quit == None: return False
        if self.expired() and o.expired(): return True
        return abs(self._time_to_quit - o._time_to_quit) < \
            timePeriodSeconds(timedelta(microseconds=1))
    def __lt__(self, o):
        try:
            if self._time_to_quit is None and o._time_to_quit is None: return False
        except Exception: pass
        if self._time_to_quit is None: return False
        if isinstance(o, timedelta):
            o = ExpiryTime(o)
        if o._time_to_quit is None: return True
        return self._time_to_quit < o._time_to_quit
    def __gt__(self, o):
        try:
            if self._time_to_quit is None and o._time_to_quit is None: return False
        except Exception: pass
        return not self.__lt__(o)
    def __le__(self, o): return self.__eq__(o) or self.__lt__(o)
    def __ge__(self, o): return self.__eq__(o) or self.__gt__(o)
    def __ne__(self, o): return not self.__eq__(o)
    def __bool__(self): return self.expired()
    def __nonzero__(self): return self.expired()

class ExpirationTimer(object):
    """Keeps track of a duration relative to an original time and
       indicates whether that duration has expired or how much time is
       left before it expires.

       May also be initialized with a duration of None, indicating
       that it should never timeout and that `remaining()` should
       return the forever value (defaulting to None).
    """
    def __init__(self, duration=None):
        self.duration = duration
        if duration is None:
            self._time_to_quit = None
        elif isinstance(duration, timedelta):
            self._time_to_quit = currentTime() + timePeriodSeconds(duration)
        else:
            self._time_to_quit = currentTime() + duration
    def expired(self):
        "Returns true if the indicated duration has passed since this was created."
        return False if self._time_to_quit is None else (currentTime() >= self._time_to_quit)
    def remaining(self, forever=None):
        """Returns a timedelta of remaining time until expiration, or 0 if the
           duration has already expired.  Returns forever if no timeout."""
        return forever if self._time_to_quit is None else \
            timedelta(seconds=self.remainingSeconds())
    def remainingSeconds(self, forever=None):
        """Similar to `remaining()`, but returns an floating point value of the
           number of remaining seconds instead of returning a
           timedelta object.
        """
        return forever if self._time_to_quit is None else \
            max(self._time_to_quit - currentTime(), 0)
    def __str__(self):
        if self._time_to_quit is None: return 'Forever'
        if self.expired():
            return 'Expired_for_%s' % \
                timedelta(seconds=currentTime() - self._time_to_quit)
        return 'Expires_in_' + str(self.remaining())
    def __eq__(self, o):
        if isinstance(o, timedelta):
            o = ExpiryTime(o)
        if self._time_to_quit == o._time_to_quit: return True
        if self._time_to_quit == None or o._time_to_quit == None: return False
        if self.expired() and o.expired(): return True
        return abs(self._time_to_quit - o._time_to_quit) < \
            timePeriodSeconds(seconds=timedelta(microseconds=1))
    def __lt__(self, o):
        try:
            if self._time_to_quit is None and o._time_to_quit is None: return False
        except Exception: pass
        if self._time_to_quit is None: return False
        if isinstance(o, timedelta):
            o = ExpiryTime(o)
        if o._time_to_quit is None: return True
        return self._time_to_quit < o._time_to_quit
    def __gt__(self, o):
        try:
            if self._time_to_quit is None and o._time_to_quit is None: return False
        except Exception: pass
        return not self.__lt__(o)
    def __le__(self, o): return self.__eq__(o) or self.__lt__(o)
    def __ge__(self, o): return self.__eq__(o) or self.__gt__(o)
    def __ne__(self, o): return not self.__eq__(o)
    def __bool__(self): return self.expired()
    def __nonzero__(self): return self.expired()
<|endoftext|>


<commit_before><commit_msg>Corrected inconsistent newlines in third_party/mesa/mesa.gyp. TEST=none BUG=none<commit_after><|endoftext|>