<commit_before>"""
  Cloud connected autonomous RC car.

  Copyright 2016 Visible Energy Inc. All Rights Reserved.
"""
__license__ = """
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.
"""
import time
import json
import serial
import string
import sys
import Queue
import threading
import subprocess

import streamer
import utils
import numpy as np
from keras.models import model_from_json

from config import DataConfig

class States:
  """ Vehicle states """
  IDLE=1
  RUNNING=2
  PAUSE=3
  STOPPED=0

class Modes:
  """ Vehicle running modes """
  AUTO=1        # fully autonomous
  TRAINING=2    # RC controlled to capture training video
  REMOTE=3      # controlled remotely through Cometa JSON/RPC

# --- Module constansts
THETA_CENTER = 90
MOTOR_NEUTRAL = 90
# filename to fetch telemetry from -- updated atomically by the controller loop at 30 Hz
TELEMFNAME = '/tmpfs/meta.txt'

# filename where to store the last frame -- used by the application loop and as parameter for the CNN prediction
FRAMEFNAME = '/tmpfs/frame.yuv'

def setup_arduino(config, logger):
  """ Arduino radio receiver and servos controller setup. """
  try:
    # set serial non-blocking
    port = serial.Serial(config['arduino']['serial'], config['arduino']['speed'], timeout=0.0, xonxoff=False, rtscts=False, dsrdtr=False)
    port.flushInput()
    port.flushOutput()
  except Exception as e:
    logger("Arduino setup: %s" % e)
    return None
  # wait the board to start
  while port.inWaiting() == 0:
    time.sleep(0.1)
  return port

class RCVehicle(object):
  """
  Vehicle controller class
    config - configuration object from config.json
    log - system logger
  """
  # current state and mode
  state = None
  mode = None

  # running options flags
  capture = False     # capturing video and telemetry for CNN training
  streaming = False   # streaming video to cloud server

  # current servo and motor values
  steering = None
  throttle = None

  # Arduino serial port
  arport = None

  # Mailbox for asynchronous commands
  mbox = Queue.Queue()

  # GPS readings
  readings={}

  # System logger
  log = None

  def __init__(self, config, logger):
    self.state=States.IDLE
    self.steering=THETA_CENTER
    self.throttle=MOTOR_NEUTRAL
    self.serial=config['serial']
    self.config = config

    # Set the system log
    self.log=logger
    self.verbose=config['app_params']['verbose']

    # CNN predicting model
    self.cnn_model = None

    self.arport=setup_arduino(config, self.log)
    while self.arport == None:
      self.log("Fatal error setting up Arduino board. Cannot proceed without properly connecting to the control board.")
      time.sleep(5)
      self.arport=setup_arduino(config, self.log)
    # TODO: remember to create a board heartbeat thread

    # Global lock
    self.glock = threading.Lock()

    # Start the main loop
    self.loop_t = threading.Thread(target=self.control_loop)
    self.loop_t.daemon = True   # force to exit on SIGINT

    self.telemetry_period=config['app_params']['telemetry_period']
    # cometa handle
    self.com = None
    return

  def load_model(self, modelpath):
    """ Load a CNN model """
    if self.state == States.AUTO:
      return False
    try:
      self.model = model_from_json(open(modelpath + ".json").read())
      self.model.load_weights(modelpath + ".h5")
    except Exception, e:
      self.log("Failed to load CNN model %s. (%s)" % (modelpath,str(e)) )
      return False

    self.log("Loaded CNN model %s." % modelpath)
    return True

  def state2run(self):
    """ State transition to RUNNING """
    if self.state == States.RUNNING:
      return
    self.steering=THETA_CENTER
    self.throttle=MOTOR_NEUTRAL
    self.state=States.RUNNING
    self.log("State RUNNING")
    return

  def state2stopped(self):
    """ State transition to PAUSE """
    if self.state == States.STOPPED:
      return
    self.steering=THETA_CENTER
    self.throttle=MOTOR_NEUTRAL
    self.state=States.STOPPED
    self.output_arduino(self.steering, self.throttle)
    self.log("State STOPPED")
    return

  def state2idle(self):
    """ State transition to IDLE """
    if self.state == States.IDLE:
      return
    self.steering=THETA_CENTER
    self.throttle=MOTOR_NEUTRAL
    self.state=States.IDLE
    self.output_arduino(self.steering, self.throttle)
    self.log("State IDLE")
    return

  def mode2auto(self):
    """ Mode transition to AUTO """
    if self.mode == Modes.AUTO:
      return
    # TODO: start the video fast video streamer

    self.mode=Modes.AUTO
    self.arport.flushInput()
    self.arport.flushOutput()
    self.log("Mode AUTO")
    return

  def mode2training(self):
    """ Mode transition to TRAINING """
    if self.mode == Modes.TRAINING:
      return
    # TODO: start the video streamer with telemetry annotations
    self.mode=Modes.TRAINING
    self.arport.flushInput()
    self.arport.flushOutput()
    self.log("Mode TRAINING")
    return

  def mode2remote(self):
    """ Mode transition to REMOTE """
    if self.mode == Modes.REMOTE:
      return

    self.arport.flushInput()
    self.arport.flushOutput()
    self.steering=THETA_CENTER
    self.throttle=MOTOR_NEUTRAL
    self.mode=Modes.REMOTE
    self.log("Mode REMOTE")
    return

  def start(self):
    """ Initial start """
    self.mode2training()
    self.state2run()
    self.loop_t.start()
    return

  def telemetry(self):
    ret = {}
    ret['type'] = 1
    ret['time'] = str(int(time.time() * 1000))
    ret['device_id'] = self.serial
    ret['state'] = self.state
    ret['mode'] = self.mode
    ret['steering'] = self.steering
    ret['throttle'] = self.throttle
    ret['GPS'] = {}
    try:
      ret['GPS']['lat'] = int(self.readings['lat'] * 10E4) / 10E4
      ret['GPS']['lon'] = int(self.readings['lon'] * 10E4) / 10E4
    except:
      # leave GPS empty
      pass
    return ret

  def input_arduino(self):
    """ Read a line composed of throttle and steering values received from the RC. """
    inputLine = ''
    if self.arport.inWaiting():
      ch = self.arport.read(1)
      while ch != b'\x0A':
        inputLine += ch
        ch = self.arport.read(1)
      try:
        # print inputLine.decode('ISO-8859-1')
        t_in, s_in = inputLine.split()
        # return the steering and throttle values from the receiver
        return int(s_in), int(t_in)
      except:
        pass
    # return current values after a reading error
    return self.steering, self.throttle

  def output_arduino(self, steering, throttle):
    """ Write steering and throttle PWM values in the [0,180] range to the controller. """
    # set steering to neutral if within an interval around 90
    steering = 90 if 88 < steering < 92 else steering
    # send a new steering PWM setting to the controller
    if self.mode == Modes.TRAINING:
      if steering != self.steering:
        self.steering = steering   # update global
        self.arport.write(('S %d\n' % self.steering).encode('ascii'))
    else:
        self.arport.write(('S %d\n' % self.steering).encode('ascii'))

    # send a new throttle PWM setting to the controller
    if self.mode == Modes.TRAINING:
      if throttle != self.throttle:
        self.throttle = throttle   # update global
        self.arport.write(('M %d\n' % self.throttle).encode('ascii'))
    else:
      self.arport.write(('M %d\n' % self.throttle).encode('ascii'))
    return

  # ---------------------------------------
  #
  def control_loop(self):
    """ Controller main loop """
    last_update=0
    steering_in=self.steering
    throttle_in=self.throttle
    mv = '/bin/mv /tmpfs/meta.tmp ' + TELEMFNAME

    # Load a CNN model -- must be done in the same thread of the prediction
    self.load_model(self.config['app_params']['model'])

    cnn_config = DataConfig()

    last_telemetry = 0.
    while True:
      now = time.time()

      # Send telemetry data
      if self.telemetry_period < now - last_telemetry:
        msg = self.telemetry()
        if self.com.send_data(json.dumps(msg)) < 0:
            self.log("Error in sending telemetry data.")
        else:
            if self.verbose:
                self.log("Sending telemetry data %s " % msg)
        last_telemetry = now

      # get inputs from RC receiver in the [0,180] range
      try:
        if self.arport.inWaiting():
          steering_in, throttle_in = self.input_arduino()
      except Exception, e:
        self.log("input_arduino: %s" % str(e))
        continue
      #
      # ------------------------------------------------------------
      #
      if self.state == States.RUNNING and self.mode == Modes.TRAINING:
        # get inputs from RC receiver in the [0,180] range
#        try:
#          if self.arport.inWaiting():
#            steering_in, throttle_in = self.input_arduino()
#        except Exception, e:
#          self.log("input_arduino: %s" % str(e))
#          continue
        # set steering to neutral if within an interval around 90
        steering_in = 90 if 87 <= steering_in < 92 else steering_in
        if self.verbose: print steering_in, throttle_in

        if self.steering == steering_in and self.throttle == throttle_in:
          # like it or not we need to sleep to avoid to hog the CPU in a spin loop
          time.sleep(0.01)
          continue

        # update telemetry file 30 times per second
        if 0.03337 < now - last_update:
          s = "%03d %03d" %  (steering_in, throttle_in)
          # create metadata file for embedding steering and throttle values in the video stream
          try:
            f = open('/tmpfs/meta.tmp', 'w', 0)
            f.write(s)
            f.close()
            # use mv that is a system call and not preempted (atomic)
            subprocess.check_call(mv, shell=True)
          except Exception, e:
            self.log("writing file: %s" % str(e))
            pass
        last_update = now
        # set new values for throttle and steering servos
        self.output_arduino(steering_in, throttle_in)
      #
      # ------------------------------------------------------------
      #
      elif self.state == States.RUNNING and self.mode == Modes.AUTO:
        # predict steering and trhottle and set the values at a rate depending on preditiction speed
        start_t = time.time()
        Y = utils.read_uyvy(FRAMEFNAME, cnn_config) # Y is of shape (1, :, :, 1) or (1, :, :, 3)
        if Y is None:
          print "image not acquired"
          continue

        # predict steering and throttle
        s, t = self.model.predict(Y[0:1])
        self.steering = np.argmax(s[0])
        self.throttle = np.argmax(t[0])

        self.steering = utils.bucket2steering(self.steering)
        self.throttle = utils.bucket2throttle(self.throttle)

        # clip the prediction for testing
        if self.throttle > 110:
          self.throttle = 110
        # insure a certain value for throttle
        if self.throttle < 97:
          self.throttle = 97

        self.output_arduino(self.steering, self.throttle)
        print self.steering, self.throttle
        dt = time.time() - start_t
        time.sleep(0.01) #0.034 - max(0.033, dt))
        print "execution time:", dt
      #
      # ------------------------------------------------------------
      #
      elif self.state == States.RUNNING and self.mode == Modes.REMOTE:
        self.output_arduino(self.steering, self.throttle)
#        time.sleep(0.2)
        time.sleep(0.01)
      #
      # ------------------------------------------------------------
      #
      elif self.state == States.IDLE:
        time.sleep(0.1)
      #
      # ------------------------------------------------------------
      #
      else:
        time.sleep(0.1)
<commit_msg>fixed typo<commit_after>"""
  Cloud connected autonomous RC car.

  Copyright 2016 Visible Energy Inc. All Rights Reserved.
"""
__license__ = """
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.
"""
import time
import json
import serial
import string
import sys
import Queue
import threading
import subprocess

import streamer
import utils
import numpy as np
from keras.models import model_from_json

from config import DataConfig

class States:
  """ Vehicle states """
  IDLE=1
  RUNNING=2
  PAUSE=3
  STOPPED=0

class Modes:
  """ Vehicle running modes """
  AUTO=1        # fully autonomous
  TRAINING=2    # RC controlled to capture training video
  REMOTE=3      # controlled remotely through Cometa JSON/RPC

# --- Module constansts
THETA_CENTER = 90
MOTOR_NEUTRAL = 90
# filename to fetch telemetry from -- updated atomically by the controller loop at 30 Hz
TELEMFNAME = '/tmpfs/meta.txt'

# filename where to store the last frame -- used by the application loop and as parameter for the CNN prediction
FRAMEFNAME = '/tmpfs/frame.yuv'

def setup_arduino(config, logger):
  """ Arduino radio receiver and servos controller setup. """
  try:
    # set serial non-blocking
    port = serial.Serial(config['arduino']['serial'], config['arduino']['speed'], timeout=0.0, xonxoff=False, rtscts=False, dsrdtr=False)
    port.flushInput()
    port.flushOutput()
  except Exception as e:
    logger("Arduino setup: %s" % e)
    return None
  # wait the board to start
  while port.inWaiting() == 0:
    time.sleep(0.1)
  return port

class RCVehicle(object):
  """
  Vehicle controller class
    config - configuration object from config.json
    log - system logger
  """
  # current state and mode
  state = None
  mode = None

  # running options flags
  capture = False     # capturing video and telemetry for CNN training
  streaming = False   # streaming video to cloud server

  # current servo and motor values
  steering = None
  throttle = None

  # Arduino serial port
  arport = None

  # Mailbox for asynchronous commands
  mbox = Queue.Queue()

  # GPS readings
  readings={}

  # System logger
  log = None

  def __init__(self, config, logger):
    self.state=States.IDLE
    self.steering=THETA_CENTER
    self.throttle=MOTOR_NEUTRAL
    self.serial=config['serial']
    self.config = config

    # Set the system log
    self.log=logger
    self.verbose=config['app_params']['verbose']

    # CNN predicting model
    self.cnn_model = None

    self.arport=setup_arduino(config, self.log)
    while self.arport == None:
      self.log("Fatal error setting up Arduino board. Cannot proceed without properly connecting to the control board.")
      time.sleep(5)
      self.arport=setup_arduino(config, self.log)
    # TODO: remember to create a board heartbeat thread

    # Global lock
    self.glock = threading.Lock()

    # Start the main loop
    self.loop_t = threading.Thread(target=self.control_loop)
    self.loop_t.daemon = True   # force to exit on SIGINT

    self.telemetry_period=config['app_params']['telemetry_period']
    # cometa handle
    self.com = None
    return

  def load_model(self, modelpath):
    """ Load a CNN model """
    if self.mode == Modes.AUTO:
      return False
    try:
      self.model = model_from_json(open(modelpath + ".json").read())
      self.model.load_weights(modelpath + ".h5")
    except Exception, e:
      self.log("Failed to load CNN model %s. (%s)" % (modelpath,str(e)) )
      return False

    self.log("Loaded CNN model %s." % modelpath)
    return True

  def state2run(self):
    """ State transition to RUNNING """
    if self.state == States.RUNNING:
      return
    self.steering=THETA_CENTER
    self.throttle=MOTOR_NEUTRAL
    self.state=States.RUNNING
    self.log("State RUNNING")
    return

  def state2stopped(self):
    """ State transition to PAUSE """
    if self.state == States.STOPPED:
      return
    self.steering=THETA_CENTER
    self.throttle=MOTOR_NEUTRAL
    self.state=States.STOPPED
    self.output_arduino(self.steering, self.throttle)
    self.log("State STOPPED")
    return

  def state2idle(self):
    """ State transition to IDLE """
    if self.state == States.IDLE:
      return
    self.steering=THETA_CENTER
    self.throttle=MOTOR_NEUTRAL
    self.state=States.IDLE
    self.output_arduino(self.steering, self.throttle)
    self.log("State IDLE")
    return

  def mode2auto(self):
    """ Mode transition to AUTO """
    if self.mode == Modes.AUTO:
      return
    # TODO: start the video fast video streamer

    self.mode=Modes.AUTO
    self.arport.flushInput()
    self.arport.flushOutput()
    self.log("Mode AUTO")
    return

  def mode2training(self):
    """ Mode transition to TRAINING """
    if self.mode == Modes.TRAINING:
      return
    # TODO: start the video streamer with telemetry annotations
    self.mode=Modes.TRAINING
    self.arport.flushInput()
    self.arport.flushOutput()
    self.log("Mode TRAINING")
    return

  def mode2remote(self):
    """ Mode transition to REMOTE """
    if self.mode == Modes.REMOTE:
      return

    self.arport.flushInput()
    self.arport.flushOutput()
    self.steering=THETA_CENTER
    self.throttle=MOTOR_NEUTRAL
    self.mode=Modes.REMOTE
    self.log("Mode REMOTE")
    return

  def start(self):
    """ Initial start """
    self.mode2training()
    self.state2run()
    self.loop_t.start()
    return

  def telemetry(self):
    ret = {}
    ret['type'] = 1
    ret['time'] = str(int(time.time() * 1000))
    ret['device_id'] = self.serial
    ret['state'] = self.state
    ret['mode'] = self.mode
    ret['steering'] = self.steering
    ret['throttle'] = self.throttle
    ret['GPS'] = {}
    try:
      ret['GPS']['lat'] = int(self.readings['lat'] * 10E4) / 10E4
      ret['GPS']['lon'] = int(self.readings['lon'] * 10E4) / 10E4
    except:
      # leave GPS empty
      pass
    return ret

  def input_arduino(self):
    """ Read a line composed of throttle and steering values received from the RC. """
    inputLine = ''
    if self.arport.inWaiting():
      ch = self.arport.read(1)
      while ch != b'\x0A':
        inputLine += ch
        ch = self.arport.read(1)
      try:
        # print inputLine.decode('ISO-8859-1')
        t_in, s_in = inputLine.split()
        # return the steering and throttle values from the receiver
        return int(s_in), int(t_in)
      except:
        pass
    # return current values after a reading error
    return self.steering, self.throttle

  def output_arduino(self, steering, throttle):
    """ Write steering and throttle PWM values in the [0,180] range to the controller. """
    # set steering to neutral if within an interval around 90
    steering = 90 if 88 < steering < 92 else steering
    # send a new steering PWM setting to the controller
    if self.mode == Modes.TRAINING:
      if steering != self.steering:
        self.steering = steering   # update global
        self.arport.write(('S %d\n' % self.steering).encode('ascii'))
    else:
        self.arport.write(('S %d\n' % self.steering).encode('ascii'))

    # send a new throttle PWM setting to the controller
    if self.mode == Modes.TRAINING:
      if throttle != self.throttle:
        self.throttle = throttle   # update global
        self.arport.write(('M %d\n' % self.throttle).encode('ascii'))
    else:
      self.arport.write(('M %d\n' % self.throttle).encode('ascii'))
    return

  # ---------------------------------------
  #
  def control_loop(self):
    """ Controller main loop """
    last_update=0
    steering_in=self.steering
    throttle_in=self.throttle
    mv = '/bin/mv /tmpfs/meta.tmp ' + TELEMFNAME

    # Load a CNN model -- must be done in the same thread of the prediction
    self.load_model(self.config['app_params']['model'])

    cnn_config = DataConfig()

    last_telemetry = 0.
    while True:
      now = time.time()

      # Send telemetry data
      if self.telemetry_period < now - last_telemetry:
        msg = self.telemetry()
        if self.com.send_data(json.dumps(msg)) < 0:
            self.log("Error in sending telemetry data.")
        else:
            if self.verbose:
                self.log("Sending telemetry data %s " % msg)
        last_telemetry = now

      # get inputs from RC receiver in the [0,180] range
      try:
        if self.arport.inWaiting():
          steering_in, throttle_in = self.input_arduino()
      except Exception, e:
        self.log("input_arduino: %s" % str(e))
        continue
      #
      # ------------------------------------------------------------
      #
      if self.state == States.RUNNING and self.mode == Modes.TRAINING:
        # get inputs from RC receiver in the [0,180] range
#        try:
#          if self.arport.inWaiting():
#            steering_in, throttle_in = self.input_arduino()
#        except Exception, e:
#          self.log("input_arduino: %s" % str(e))
#          continue
        # set steering to neutral if within an interval around 90
        steering_in = 90 if 87 <= steering_in < 92 else steering_in
        if self.verbose: print steering_in, throttle_in

        if self.steering == steering_in and self.throttle == throttle_in:
          # like it or not we need to sleep to avoid to hog the CPU in a spin loop
          time.sleep(0.01)
          continue

        # update telemetry file 30 times per second
        if 0.03337 < now - last_update:
          s = "%03d %03d" %  (steering_in, throttle_in)
          # create metadata file for embedding steering and throttle values in the video stream
          try:
            f = open('/tmpfs/meta.tmp', 'w', 0)
            f.write(s)
            f.close()
            # use mv that is a system call and not preempted (atomic)
            subprocess.check_call(mv, shell=True)
          except Exception, e:
            self.log("writing file: %s" % str(e))
            pass
        last_update = now
        # set new values for throttle and steering servos
        self.output_arduino(steering_in, throttle_in)
      #
      # ------------------------------------------------------------
      #
      elif self.state == States.RUNNING and self.mode == Modes.AUTO:
        # predict steering and trhottle and set the values at a rate depending on preditiction speed
        start_t = time.time()
        Y = utils.read_uyvy(FRAMEFNAME, cnn_config) # Y is of shape (1, :, :, 1) or (1, :, :, 3)
        if Y is None:
          print "image not acquired"
          continue

        # predict steering and throttle
        s, t = self.model.predict(Y[0:1])
        self.steering = np.argmax(s[0])
        self.throttle = np.argmax(t[0])

        self.steering = utils.bucket2steering(self.steering)
        self.throttle = utils.bucket2throttle(self.throttle)

        # clip the prediction for testing
        if self.throttle > 110:
          self.throttle = 110
        # insure a certain value for throttle
        if self.throttle < 97:
          self.throttle = 97

        self.output_arduino(self.steering, self.throttle)
        print self.steering, self.throttle
        dt = time.time() - start_t
        time.sleep(0.01) #0.034 - max(0.033, dt))
        print "execution time:", dt
      #
      # ------------------------------------------------------------
      #
      elif self.state == States.RUNNING and self.mode == Modes.REMOTE:
        self.output_arduino(self.steering, self.throttle)
#        time.sleep(0.2)
        time.sleep(0.01)
      #
      # ------------------------------------------------------------
      #
      elif self.state == States.IDLE:
        time.sleep(0.1)
      #
      # ------------------------------------------------------------
      #
      else:
        time.sleep(0.1)
<|endoftext|>


<commit_before>"""
Support the OwnTracks platform.

For more details about this platform, please refer to the documentation at
https://home-assistant.io/components/device_tracker.owntracks/
"""
import json
import logging
import threading
from collections import defaultdict

import homeassistant.components.mqtt as mqtt
from homeassistant.const import STATE_HOME

DEPENDENCIES = ['mqtt']

REGIONS_ENTERED = defaultdict(list)
MOBILE_BEACONS_ACTIVE = defaultdict(list)

BEACON_DEV_ID = 'beacon'

LOCATION_TOPIC = 'owntracks/+/+'
EVENT_TOPIC = 'owntracks/+/+/event'

_LOGGER = logging.getLogger(__name__)

LOCK = threading.Lock()

CONF_MAX_GPS_ACCURACY = 'max_gps_accuracy'

def setup_scanner(hass, config, see):
    """Setup an OwnTracks tracker."""
    max_gps_accuracy = config.get(CONF_MAX_GPS_ACCURACY)

    def owntracks_location_update(topic, payload, qos):
        """MQTT message received."""
        # Docs on available data:
        # http://owntracks.org/booklet/tech/json/#_typelocation
        try:
            data = json.loads(payload)
        except ValueError:
            # If invalid JSON
            _LOGGER.error(
                'Unable to parse payload as JSON: %s', payload)
            return

        if (not isinstance(data, dict) or data.get('_type') != 'location') or (
                'acc' in data and max_gps_accuracy is not None and data[
                    'acc'] > max_gps_accuracy):
            return

        dev_id, kwargs = _parse_see_args(topic, data)

        # Block updates if we're in a region
        with LOCK:
            if REGIONS_ENTERED[dev_id]:
                _LOGGER.debug(
                    "location update ignored - inside region %s",
                    REGIONS_ENTERED[-1])
                return

            see(**kwargs)
            see_beacons(dev_id, kwargs)

    def owntracks_event_update(topic, payload, qos):
        # pylint: disable=too-many-branches, too-many-statements
        """MQTT event (geofences) received."""
        # Docs on available data:
        # http://owntracks.org/booklet/tech/json/#_typetransition
        try:
            data = json.loads(payload)
        except ValueError:
            # If invalid JSON
            _LOGGER.error(
                'Unable to parse payload as JSON: %s', payload)
            return

        if not isinstance(data, dict) or data.get('_type') != 'transition':
            return

        # OwnTracks uses - at the start of a beacon zone
        # to switch on 'hold mode' - ignore this
        location = data['desc'].lstrip("-")
        if location.lower() == 'home':
            location = STATE_HOME

        dev_id, kwargs = _parse_see_args(topic, data)

        if data['event'] == 'enter':
            zone = hass.states.get("zone.{}".format(location))
            with LOCK:
                if zone is None:
                    if data['t'] == 'b':
                        # Not a HA zone, and a beacon so assume mobile
                        beacons = MOBILE_BEACONS_ACTIVE[dev_id]
                        if location not in beacons:
                            beacons.append(location)
                        _LOGGER.info("Added beacon %s", location)
                else:
                    # Normal region
                    regions = REGIONS_ENTERED[dev_id]
                    if location not in regions:
                        regions.append(location)
                    _LOGGER.info("Enter region %s", location)
                    _set_gps_from_zone(kwargs, location, zone)

                see(**kwargs)
                see_beacons(dev_id, kwargs)

        elif data['event'] == 'leave':
            with LOCK:
                regions = REGIONS_ENTERED[dev_id]
                if location in regions:
                    regions.remove(location)
                new_region = regions[-1] if regions else None

                if new_region:
                    # Exit to previous region
                    zone = hass.states.get("zone.{}".format(new_region))
                    _set_gps_from_zone(kwargs, new_region, zone)
                    _LOGGER.info("Exit to %s", new_region)
                    see(**kwargs)
                    see_beacons(dev_id, kwargs)

                else:
                    _LOGGER.info("Exit to GPS")
                    # Check for GPS accuracy
                    if not ('acc' in data and
                            max_gps_accuracy is not None and
                            data['acc'] > max_gps_accuracy):

                        see(**kwargs)
                        see_beacons(dev_id, kwargs)
                    else:
                        _LOGGER.info("Inaccurate GPS reported")

                beacons = MOBILE_BEACONS_ACTIVE[dev_id]
                if location in beacons:
                    beacons.remove(location)
                    _LOGGER.info("Remove beacon %s", location)

        else:
            _LOGGER.error(
                'Misformatted mqtt msgs, _type=transition, event=%s',
                data['event'])
            return

    def see_beacons(dev_id, kwargs_param):
        """Set active beacons to the current location."""
        kwargs = kwargs_param.copy()
        # the battery state applies to the tracking device, not the beacon
        kwargs.pop('battery', None)
        for beacon in MOBILE_BEACONS_ACTIVE[dev_id]:
            kwargs['dev_id'] = "{}_{}".format(BEACON_DEV_ID, beacon)
            kwargs['host_name'] = beacon
            see(**kwargs)

    mqtt.subscribe(hass, LOCATION_TOPIC, owntracks_location_update, 1)
    mqtt.subscribe(hass, EVENT_TOPIC, owntracks_event_update, 1)

    return True

def _parse_see_args(topic, data):
    """Parse the OwnTracks location parameters, into the format see expects."""
    parts = topic.split('/')
    dev_id = '{}_{}'.format(parts[1], parts[2])
    host_name = parts[1]
    kwargs = {
        'dev_id': dev_id,
        'host_name': host_name,
        'gps': (data['lat'], data['lon'])
    }
    if 'acc' in data:
        kwargs['gps_accuracy'] = data['acc']
    if 'batt' in data:
        kwargs['battery'] = data['batt']
    return dev_id, kwargs

def _set_gps_from_zone(kwargs, location, zone):
    """Set the see parameters from the zone parameters."""
    if zone is not None:
        kwargs['gps'] = (
            zone.attributes['latitude'],
            zone.attributes['longitude'])
        kwargs['gps_accuracy'] = zone.attributes['radius']
        kwargs['location_name'] = location
    return kwargs
<commit_msg>Cast acc to float before comparison.<commit_after>"""
Support the OwnTracks platform.

For more details about this platform, please refer to the documentation at
https://home-assistant.io/components/device_tracker.owntracks/
"""
import json
import logging
import threading
from collections import defaultdict

import homeassistant.components.mqtt as mqtt
from homeassistant.const import STATE_HOME
from homeassistant.util import convert

DEPENDENCIES = ['mqtt']

REGIONS_ENTERED = defaultdict(list)
MOBILE_BEACONS_ACTIVE = defaultdict(list)

BEACON_DEV_ID = 'beacon'

LOCATION_TOPIC = 'owntracks/+/+'
EVENT_TOPIC = 'owntracks/+/+/event'

_LOGGER = logging.getLogger(__name__)

LOCK = threading.Lock()

CONF_MAX_GPS_ACCURACY = 'max_gps_accuracy'

def setup_scanner(hass, config, see):
    """Setup an OwnTracks tracker."""
    max_gps_accuracy = config.get(CONF_MAX_GPS_ACCURACY)

    def owntracks_location_update(topic, payload, qos):
        """MQTT message received."""
        # Docs on available data:
        # http://owntracks.org/booklet/tech/json/#_typelocation
        try:
            data = json.loads(payload)
        except ValueError:
            # If invalid JSON
            _LOGGER.error(
                'Unable to parse payload as JSON: %s', payload)
            return

        if (not isinstance(data, dict) or data.get('_type') != 'location') or (
                max_gps_accuracy is not None and
                convert(data.get('acc'), float, 0.0) > max_gps_accuracy):
            return

        dev_id, kwargs = _parse_see_args(topic, data)

        # Block updates if we're in a region
        with LOCK:
            if REGIONS_ENTERED[dev_id]:
                _LOGGER.debug(
                    "location update ignored - inside region %s",
                    REGIONS_ENTERED[-1])
                return

            see(**kwargs)
            see_beacons(dev_id, kwargs)

    def owntracks_event_update(topic, payload, qos):
        # pylint: disable=too-many-branches, too-many-statements
        """MQTT event (geofences) received."""
        # Docs on available data:
        # http://owntracks.org/booklet/tech/json/#_typetransition
        try:
            data = json.loads(payload)
        except ValueError:
            # If invalid JSON
            _LOGGER.error(
                'Unable to parse payload as JSON: %s', payload)
            return

        if not isinstance(data, dict) or data.get('_type') != 'transition':
            return

        # OwnTracks uses - at the start of a beacon zone
        # to switch on 'hold mode' - ignore this
        location = data['desc'].lstrip("-")
        if location.lower() == 'home':
            location = STATE_HOME

        dev_id, kwargs = _parse_see_args(topic, data)

        if data['event'] == 'enter':
            zone = hass.states.get("zone.{}".format(location))
            with LOCK:
                if zone is None:
                    if data['t'] == 'b':
                        # Not a HA zone, and a beacon so assume mobile
                        beacons = MOBILE_BEACONS_ACTIVE[dev_id]
                        if location not in beacons:
                            beacons.append(location)
                        _LOGGER.info("Added beacon %s", location)
                else:
                    # Normal region
                    regions = REGIONS_ENTERED[dev_id]
                    if location not in regions:
                        regions.append(location)
                    _LOGGER.info("Enter region %s", location)
                    _set_gps_from_zone(kwargs, location, zone)

                see(**kwargs)
                see_beacons(dev_id, kwargs)

        elif data['event'] == 'leave':
            with LOCK:
                regions = REGIONS_ENTERED[dev_id]
                if location in regions:
                    regions.remove(location)
                new_region = regions[-1] if regions else None

                if new_region:
                    # Exit to previous region
                    zone = hass.states.get("zone.{}".format(new_region))
                    _set_gps_from_zone(kwargs, new_region, zone)
                    _LOGGER.info("Exit to %s", new_region)
                    see(**kwargs)
                    see_beacons(dev_id, kwargs)

                else:
                    _LOGGER.info("Exit to GPS")
                    # Check for GPS accuracy
                    if not ('acc' in data and
                            max_gps_accuracy is not None and
                            data['acc'] > max_gps_accuracy):

                        see(**kwargs)
                        see_beacons(dev_id, kwargs)
                    else:
                        _LOGGER.info("Inaccurate GPS reported")

                beacons = MOBILE_BEACONS_ACTIVE[dev_id]
                if location in beacons:
                    beacons.remove(location)
                    _LOGGER.info("Remove beacon %s", location)

        else:
            _LOGGER.error(
                'Misformatted mqtt msgs, _type=transition, event=%s',
                data['event'])
            return

    def see_beacons(dev_id, kwargs_param):
        """Set active beacons to the current location."""
        kwargs = kwargs_param.copy()
        # the battery state applies to the tracking device, not the beacon
        kwargs.pop('battery', None)
        for beacon in MOBILE_BEACONS_ACTIVE[dev_id]:
            kwargs['dev_id'] = "{}_{}".format(BEACON_DEV_ID, beacon)
            kwargs['host_name'] = beacon
            see(**kwargs)

    mqtt.subscribe(hass, LOCATION_TOPIC, owntracks_location_update, 1)
    mqtt.subscribe(hass, EVENT_TOPIC, owntracks_event_update, 1)

    return True

def _parse_see_args(topic, data):
    """Parse the OwnTracks location parameters, into the format see expects."""
    parts = topic.split('/')
    dev_id = '{}_{}'.format(parts[1], parts[2])
    host_name = parts[1]
    kwargs = {
        'dev_id': dev_id,
        'host_name': host_name,
        'gps': (data['lat'], data['lon'])
    }
    if 'acc' in data:
        kwargs['gps_accuracy'] = data['acc']
    if 'batt' in data:
        kwargs['battery'] = data['batt']
    return dev_id, kwargs

def _set_gps_from_zone(kwargs, location, zone):
    """Set the see parameters from the zone parameters."""
    if zone is not None:
        kwargs['gps'] = (
            zone.attributes['latitude'],
            zone.attributes['longitude'])
        kwargs['gps_accuracy'] = zone.attributes['radius']
        kwargs['location_name'] = location
    return kwargs
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
##
## $Id$
##
## This file is part of CDS Invenio.
## Copyright (C) 2002, 2003, 2004, 2005, 2006, 2007 CERN.
##
## CDS Invenio is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License as
## published by the Free Software Foundation; either version 2 of the
## License, or (at your option) any later version.
##
## CDS Invenio is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
## General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with CDS Invenio; if not, write to the Free Software Foundation, Inc.,
## 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.

"""Testing functions for the page getter module.
"""

__revision__ = "$Id$"

import unittest

from invenio.websearch_external_collections_getter import HTTPAsyncPageGetter, async_download

class AsyncDownloadTest(unittest.TestCase):
    """Test suite for websearch_external_collections_*"""

    def test_async_download(self):
        """websearch_external_collections_getter - asynchronous download"""

        ## Test varius cases for the async_download function:
        ##   - test 2 workings pages : google, kernel.org
        ##   - test 1 unresolvable name : rjfreijoiregjreoijgoirg.fr
        ##   - test 1 bad ip : 1.2.3.4
        ## Return the list of errors.
        # {'url': 'http://public.web.cern.ch/public/', 'content': "<title>CERN - The world's largest particle physics laboratory</title>"},
        checks = [
            {'url': 'http://cdsware.cern.ch/invenio/index.html', 'content': '<title>CDS Invenio: Overview</title>'},
                    {'url': 'http://cdsware.cern.ch/invenio/index.html', 'content': '<title>CDS Invenio: Overview</title>'},
                    {'url': 'http://rjfreijoiregjreoijgoirg.fr'},
                    {'url': 'http://1.2.3.4/'} ]

        def finished(pagegetter, check, current_time):
            """Function called when a page is received."""
            is_ok = pagegetter.status is not None

            if check.has_key('content') and is_ok:
                is_ok = pagegetter.data.find(check['content']) > 0

            check['result'] = is_ok == check.has_key('content')

        pagegetters = [HTTPAsyncPageGetter(check['url']) for check in checks]
        finished_list = async_download(pagegetters, finished, checks, 20)

        for (finished, check) in zip(finished_list, checks):
            if not finished:
                check['result'] = not check.has_key('content')

        errors = [check for check in checks if not check['result']]

        self.assertEqual(errors, [])

def create_test_suite():
    """Return test suite for the external collection tests."""
    return unittest.TestSuite((unittest.makeSuite(AsyncDownloadTest, 'test'),))

if __name__ == "__main__":
    unittest.TextTestRunner(verbosity=2).run(create_test_suite())

<commit_msg>Clarified documentation.<commit_after># -*- coding: utf-8 -*-
##
## $Id$
##
## This file is part of CDS Invenio.
## Copyright (C) 2002, 2003, 2004, 2005, 2006, 2007 CERN.
##
## CDS Invenio is free software; you can redistribute it and/or
## modify it under the terms of the GNU General Public License as
## published by the Free Software Foundation; either version 2 of the
## License, or (at your option) any later version.
##
## CDS Invenio is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
## General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with CDS Invenio; if not, write to the Free Software Foundation, Inc.,
## 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.

"""Testing functions for the page getter module.
"""

__revision__ = "$Id$"

import unittest

from invenio.websearch_external_collections_getter import HTTPAsyncPageGetter, async_download

class AsyncDownloadTest(unittest.TestCase):
    """Test suite for websearch_external_collections_*"""

    def test_async_download(self):
        """websearch_external_collections_getter - asynchronous download"""

        ## Test various cases for the async_download function:
        ##   - test 1 working page: cdsware
        ##   - test 1 unresolvable name: rjfreijoiregjreoijgoirg.fr
        ##   - test 1 bad IP: 1.2.3.4
        ## Return the list of errors.
        checks = [
            {'url': 'http://cdsware.cern.ch/invenio/index.html', 'content': '<title>CDS Invenio: Overview</title>'},
            {'url': 'http://rjfreijoiregjreoijgoirg.fr'},
            {'url': 'http://1.2.3.4/'} ]

        def finished(pagegetter, check, current_time):
            """Function called when a page is received."""
            is_ok = pagegetter.status is not None

            if check.has_key('content') and is_ok:
                is_ok = pagegetter.data.find(check['content']) > 0

            check['result'] = is_ok == check.has_key('content')

        pagegetters = [HTTPAsyncPageGetter(check['url']) for check in checks]
        finished_list = async_download(pagegetters, finished, checks, 20)

        for (finished, check) in zip(finished_list, checks):
            if not finished:
                check['result'] = not check.has_key('content')

        errors = [check for check in checks if not check['result']]

        self.assertEqual(errors, [])

def create_test_suite():
    """Return test suite for the external collection tests."""
    return unittest.TestSuite((unittest.makeSuite(AsyncDownloadTest, 'test'),))

if __name__ == "__main__":
    unittest.TextTestRunner(verbosity=2).run(create_test_suite())

<|endoftext|>


<commit_before>#!/usr/bin/env python
# vim: ts=4 sts=4 et sw=4

# A script/bot for reddit to manage user awarded karma
#
# see karmaflair.ini to set options

from praw import Reddit
from praw import helpers
from reddit import reddit_auth
from reddit import reddit_reply_to_comment
import ConfigParser
from string import Template
import psycopg2
import sys
import re
import time
from datetime import datetime

# globals
debug_level = ''
cfg_file = None
r = None
conn = None
cur = None

# helper class for tee logging
class FlushOutput(object):
    def __init__(self):
        self.terminal = sys.stdout

    def write(self, message):
        self.terminal.write(message)
        self.terminal.flush()

# sys.stdout = FlushOutput()

def get_reply_text(reply_type, reply_vars):
    template = open('tmpl/karmaflair/' + reply_type + '.tpl')

    return Template(template.read()).substitute(reply_vars)

def check_for_reply(submission, name, granter):
    replied = False

    try:
        # get their total karma from the db
        cur.execute("SELECT replied FROM karma WHERE id=%s AND name=%s AND granter=%s AND replied=TRUE", (submission.id, name, granter,))
        result = cur.fetchone()

        if result is not None:
            replied = True

            if debug_level == 'DEBUG':
                print('[{}] [DEBUG] Reply exists for submission {}'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), submission.id))
    except Exception as e:
        conn.rollback()

        sys.stderr.write('[{}] [ERROR]: {}\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), e))
    else:
        conn.commit()

        return replied

def grant_karma(comment, submission, name, granter, reply_vars):
    try:
        cur.execute("INSERT INTO " + cfg_file.get('karmaflair', 'dbtablename') + " (id, name, granter) VALUES (%s, %s, %s)",
                (submission.id, name, granter))
    except psycopg2.IntegrityError as e:
        conn.rollback()

        if e.pgcode == '23505':
            # unique key violation, already granted
            if debug_level == 'NOTICE' or debug_level == 'DEBUG':
                print('[{}] [NOTICE] Karma has already been granted to {} by {}, for submission {}'
                        .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), name, granter, submission.id))

            if check_for_reply(submission, name, granter):
                reddit_reply_to_comment(comment, get_reply_text('already_awarded', reply_vars))
        else:
            sys.stderr.write('[{}] [ERROR]: {}\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), e))
    else:
        conn.commit()

        print('[{}] Karma successfully granted to {} by {}, for submission {}'
                .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), name, granter, submission.id))

        # reply and update karma flair
        if check_for_reply(submission, name, granter):
            reddit_reply_to_comment(comment, get_reply_text('successful_award', reply_vars))

        set_karma_flair(name)

def set_karma_flair(name):
    try:
        # get their total karma from the db
        cur.execute("SELECT name, count(*) AS karma FROM karma WHERE name=%s GROUP BY name", (name,))
        result = cur.fetchone()

        if result[1]:
            # grab their existing flair info
            subreddit = cfg_file.get('karmaflair', 'subreddit')
            current_flair = r.get_flair(subreddit, name)
            karma_flair_text = str(result[1]) + " Karma"

            r.set_flair(subreddit, name, karma_flair_text, current_flair['flair_css_class'])
    except Exception as e:
        conn.rollback()

        sys.stderr.write('[{}] [ERROR]: {}\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), e))
    else:
        conn.commit()

        if debug_level == 'NOTICE' or debug_level == 'DEBUG':
            print('[{}] [NOTICE] Karma flair successfully updated for {} to {}'
                    .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), name, karma_flair_text))

def process_comment_command(command, command_type, valid_commands, comment, submission, parent=None):
    if command == 'karma' and command_type == '+' and parent is not None:
        # dict of vars for template completion
        reply_vars = {
            'name': comment.author.name,
            'parent_name': parent.author.name,
        }

        # valid grant karma command, check for additional criteria
        while True:
            # request must have correct link flair
            if submission.link_flair_text != cfg_file.get('karmaflair', 'valid_flair_text'):
                reddit_reply_to_comment(comment, get_reply_text('invalid_link_flair', reply_vars))
                break

            # user cannot grant karma to themselves
            if parent.author.name == comment.author.name:
                reddit_reply_to_comment(comment, get_reply_text('award_to_self', reply_vars))
                break

            # user granting karma must be the same as the submitter, or the parent must be the submitter
            if comment.author.name != comment.link_author and parent.author.name != comment.link_author:
                reddit_reply_to_comment(comment, get_reply_text('invalid_author', reply_vars))
                break

            # cannot grant karma to another command
            if re.match("^([\+|-])(" + valid_commands + ")$", parent.body.lower().strip()):
                reddit_reply_to_comment(comment, get_reply_text('award_to_command', reply_vars))
                break

            grant_karma(comment, submission, parent.author.name, comment.author.name, reply_vars)

            break

def main():
    global cfg_file
    global debug_level
    global r
    global conn
    global cur

    # read ini and set config
    cfg_file = ConfigParser.RawConfigParser()
    while True:
        try:
            cfg_file.read('karmaflair.ini')
            break
        except Exception as e:
            sys.stderr.write('[{}] [ERROR]: {}\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), e))
            sys.exit()

    debug_level = cfg_file.get('debug', 'level')
    mode = cfg_file.get('general', 'mode')
    loop_time = cfg_file.getint('general', 'loop_time')
    subreddit = cfg_file.get('karmaflair', 'subreddit')
    valid_commands = cfg_file.get('karmaflair', 'valid_commands')

    # main loop at set interval if mode is set to 'continuous'
    while True:
        print('[{}] Starting karma flair...'
              .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S')))

        try:
            # connect to db
            conn = psycopg2.connect('dbname=' + cfg_file.get('karmaflair', 'dbname') + ' user=' + cfg_file.get('karmaflair', 'dbuser'))
            cur = conn.cursor()

            # login
            r = Reddit(user_agent=cfg_file.get('auth', 'user_agent'))
            reddit_auth(r, cfg_file, debug_level)

            # retrieve comments, stream will go back limit # of comments from start
            for comment in helpers.comment_stream(r, subreddit, limit=100, verbosity=0):
                if not comment.is_root:
                    # comment has a parent that isn't the submission, so is a candidate for a command
                    if debug_level == 'DEBUG':
                        print('[{}] [DEBUG] Checking comment posted at {} by {}'
                            .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), datetime.utcfromtimestamp(comment.created_utc), comment.author.name))

                    # command format is to start with a + or - and be the only text in the comment, whitespace is excluded
                    match = re.match("^([\+|-])(" + valid_commands + ")$", comment.body.lower().strip())

                    if match is not None and match.group(2):
                        # comment is both a sub comment and contains a valid command
                        command = match.group(2)
                        command_type = match.group(1)

                        submission = r.get_info(thing_id=comment.link_id)
                        parent = r.get_info(thing_id=comment.parent_id)

                        if debug_level == 'DEBUG':
                            print('[{}] [DEBUG] Processing comment command: {}{}'
                                .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), command_type, command))

                        process_comment_command(command, command_type, valid_commands, comment, submission, parent)

            if mode == 'continuous':
                print('[{}] Pausing karma flair...'
                    .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S')))
                time.sleep(loop_time)
            else:
                cur.close()
                conn.close()
                break
        except KeyboardInterrupt, SystemExit:
            cur.close()
            conn.close()

            print('[{}] Stopping summon karma...'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S')))
            break
        except Exception as e:
            sys.stderr.write('[{}] [ERROR]: {}\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), e))

            if mode == 'continuous':
                time.sleep(loop_time)
            else:
                break

if __name__ == '__main__':
    main()
<commit_msg>enable output flushing<commit_after>#!/usr/bin/env python
# vim: ts=4 sts=4 et sw=4

# A script/bot for reddit to manage user awarded karma
#
# see karmaflair.ini to set options

from praw import Reddit
from praw import helpers
from reddit import reddit_auth
from reddit import reddit_reply_to_comment
import ConfigParser
from string import Template
import psycopg2
import sys
import re
import time
from datetime import datetime

# globals
debug_level = ''
cfg_file = None
r = None
conn = None
cur = None

# helper class for tee logging
class FlushOutput(object):
    def __init__(self):
        self.terminal = sys.stdout

    def write(self, message):
        self.terminal.write(message)
        self.terminal.flush()

sys.stdout = FlushOutput()

def get_reply_text(reply_type, reply_vars):
    template = open('tmpl/karmaflair/' + reply_type + '.tpl')

    return Template(template.read()).substitute(reply_vars)

def check_for_reply(submission, name, granter):
    replied = False

    try:
        # get their total karma from the db
        cur.execute("SELECT replied FROM karma WHERE id=%s AND name=%s AND granter=%s AND replied=TRUE", (submission.id, name, granter,))
        result = cur.fetchone()

        if result is not None:
            replied = True

            if debug_level == 'DEBUG':
                print('[{}] [DEBUG] Reply exists for submission {}'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), submission.id))
    except Exception as e:
        conn.rollback()

        sys.stderr.write('[{}] [ERROR]: {}\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), e))
        sys.stderr.flush()
    else:
        conn.commit()

        return replied

def grant_karma(comment, submission, name, granter, reply_vars):
    try:
        cur.execute("INSERT INTO " + cfg_file.get('karmaflair', 'dbtablename') + " (id, name, granter) VALUES (%s, %s, %s)",
                (submission.id, name, granter))
    except psycopg2.IntegrityError as e:
        conn.rollback()

        if e.pgcode == '23505':
            # unique key violation, already granted
            if debug_level == 'NOTICE' or debug_level == 'DEBUG':
                print('[{}] [NOTICE] Karma has already been granted to {} by {}, for submission {}'
                        .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), name, granter, submission.id))

            if check_for_reply(submission, name, granter):
                reddit_reply_to_comment(comment, get_reply_text('already_awarded', reply_vars))
        else:
            sys.stderr.write('[{}] [ERROR]: {}\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), e))
            sys.stderr.flush()
    else:
        conn.commit()

        print('[{}] Karma successfully granted to {} by {}, for submission {}'
                .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), name, granter, submission.id))

        # reply and update karma flair
        if check_for_reply(submission, name, granter):
            reddit_reply_to_comment(comment, get_reply_text('successful_award', reply_vars))

        set_karma_flair(name)

def set_karma_flair(name):
    try:
        # get their total karma from the db
        cur.execute("SELECT name, count(*) AS karma FROM karma WHERE name=%s GROUP BY name", (name,))
        result = cur.fetchone()

        if result[1]:
            # grab their existing flair info
            subreddit = cfg_file.get('karmaflair', 'subreddit')
            current_flair = r.get_flair(subreddit, name)
            karma_flair_text = str(result[1]) + " Karma"

            r.set_flair(subreddit, name, karma_flair_text, current_flair['flair_css_class'])
    except Exception as e:
        conn.rollback()

        sys.stderr.write('[{}] [ERROR]: {}\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), e))
        sys.stderr.flush()
    else:
        conn.commit()

        if debug_level == 'NOTICE' or debug_level == 'DEBUG':
            print('[{}] [NOTICE] Karma flair successfully updated for {} to {}'
                    .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), name, karma_flair_text))

def process_comment_command(command, command_type, valid_commands, comment, submission, parent=None):
    if command == 'karma' and command_type == '+' and parent is not None:
        # dict of vars for template completion
        reply_vars = {
            'name': comment.author.name,
            'parent_name': parent.author.name,
        }

        # valid grant karma command, check for additional criteria
        while True:
            # request must have correct link flair
            if submission.link_flair_text != cfg_file.get('karmaflair', 'valid_flair_text'):
                reddit_reply_to_comment(comment, get_reply_text('invalid_link_flair', reply_vars))
                break

            # user cannot grant karma to themselves
            if parent.author.name == comment.author.name:
                reddit_reply_to_comment(comment, get_reply_text('award_to_self', reply_vars))
                break

            # user granting karma must be the same as the submitter, or the parent must be the submitter
            if comment.author.name != comment.link_author and parent.author.name != comment.link_author:
                reddit_reply_to_comment(comment, get_reply_text('invalid_author', reply_vars))
                break

            # cannot grant karma to another command
            if re.match("^([\+|-])(" + valid_commands + ")$", parent.body.lower().strip()):
                reddit_reply_to_comment(comment, get_reply_text('award_to_command', reply_vars))
                break

            grant_karma(comment, submission, parent.author.name, comment.author.name, reply_vars)

            break

def main():
    global cfg_file
    global debug_level
    global r
    global conn
    global cur

    # read ini and set config
    cfg_file = ConfigParser.RawConfigParser()
    while True:
        try:
            cfg_file.read('karmaflair.ini')
            break
        except Exception as e:
            sys.stderr.write('[{}] [ERROR]: {}\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), e))
            sys.stderr.flush()
            sys.exit()

    debug_level = cfg_file.get('debug', 'level')
    mode = cfg_file.get('general', 'mode')
    loop_time = cfg_file.getint('general', 'loop_time')
    subreddit = cfg_file.get('karmaflair', 'subreddit')
    valid_commands = cfg_file.get('karmaflair', 'valid_commands')

    # main loop at set interval if mode is set to 'continuous'
    while True:
        print('[{}] Starting karma flair...'
              .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S')))

        try:
            # connect to db
            conn = psycopg2.connect('dbname=' + cfg_file.get('karmaflair', 'dbname') + ' user=' + cfg_file.get('karmaflair', 'dbuser'))
            cur = conn.cursor()

            # login
            r = Reddit(user_agent=cfg_file.get('auth', 'user_agent'))
            reddit_auth(r, cfg_file, debug_level)

            # retrieve comments, stream will go back limit # of comments from start
            for comment in helpers.comment_stream(r, subreddit, limit=100, verbosity=0):
                if not comment.is_root:
                    # comment has a parent that isn't the submission, so is a candidate for a command
                    if debug_level == 'DEBUG':
                        print('[{}] [DEBUG] Checking comment posted at {} by {}'
                            .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), datetime.utcfromtimestamp(comment.created_utc), comment.author.name))

                    # command format is to start with a + or - and be the only text in the comment, whitespace is excluded
                    match = re.match("^([\+|-])(" + valid_commands + ")$", comment.body.lower().strip())

                    if match is not None and match.group(2):
                        # comment is both a sub comment and contains a valid command
                        command = match.group(2)
                        command_type = match.group(1)

                        submission = r.get_info(thing_id=comment.link_id)
                        parent = r.get_info(thing_id=comment.parent_id)

                        if debug_level == 'DEBUG':
                            print('[{}] [DEBUG] Processing comment command: {}{}'
                                .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), command_type, command))

                        process_comment_command(command, command_type, valid_commands, comment, submission, parent)

            if mode == 'continuous':
                print('[{}] Pausing karma flair...'
                    .format(datetime.now().strftime('%Y-%m-%d %H:%M:%S')))
                time.sleep(loop_time)
            else:
                cur.close()
                conn.close()
                break
        except KeyboardInterrupt, SystemExit:
            cur.close()
            conn.close()

            print('[{}] Stopping summon karma...'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S')))
            break
        except Exception as e:
            sys.stderr.write('[{}] [ERROR]: {}\n'.format(datetime.now().strftime('%Y-%m-%d %H:%M:%S'), e))
            sys.stderr.flush()

            if mode == 'continuous':
                time.sleep(loop_time)
            else:
                break

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>from __future__ import print_function
import tensorflow as tf
import numpy as np

import argparse
import os
import json
import collections

from bayou.core.model import Model
from bayou.core.utils import CHILD_EDGE, SIBLING_EDGE
from bayou.core.utils import read_config

MAX_GEN_UNTIL_STOP = 20

def infer(clargs):
    with tf.Session() as sess:
        predictor = BayesianPredictor(clargs.save, sess)
        err = 0
        asts = []
        if clargs.evidence_file:
            with open(clargs.evidence_file) as f:
                js = json.load(f)
        for c in range(clargs.n):
            print('Generated {} ASTs ({} errors)'.format(c, err), end='\r')
            try:
                if clargs.random:
                    psi = predictor.psi_random()
                else:
                    psi = predictor.psi_from_evidence(js)
                ast = predictor.generate_ast(psi)
                if clargs.plot2d:
                    ast['psi'] = list(psi[0])
                asts.append(ast)
                c += 1
            except AssertionError:
                err += 1
    if clargs.plot2d:
        if predictor.model.config.latent_size == 2:
            plot2d(asts)
        else:
            print('Latent space is not 2-dimensional.. cannot plot')

    if clargs.output_file is None:
        print(json.dumps({'asts': asts}, indent=2))
    else:
        with open(clargs.output_file, 'w') as f:
            json.dump({'asts': asts}, fp=f, indent=2)
    print('Number of errors: {}'.format(err))

class BayesianPredictor(object):

    def __init__(self, save, sess):
        self.sess = sess

        # load the saved config
        with open(os.path.join(save, 'config.json')) as f:
            config = read_config(json.load(f), save_dir=save, infer=True)
        self.model = Model(config, True)

        # restore the saved model
        tf.global_variables_initializer().run()
        saver = tf.train.Saver(tf.global_variables())
        ckpt = tf.train.get_checkpoint_state(save)
        saver.restore(self.sess, ckpt.model_checkpoint_path)

    def infer(self, evidences):
        psi = self.psi_from_evidence(evidences)
        return self.generate_ast(psi)

    def psi_random(self):
        return np.random.normal(size=[1, self.model.config.latent_size])

    def psi_from_evidence(self, js_evidences):
        return self.model.infer_psi(self.sess, js_evidences)

    def gen_until_STOP(self, psi, in_nodes, in_edges, check_call=False):
        ast = []
        nodes, edges = in_nodes[:], in_edges[:]
        num = 0
        while True:
            assert num < MAX_GEN_UNTIL_STOP # exception caught in main
            dist = self.model.infer_ast(self.sess, psi, nodes, edges)
            idx = np.random.choice(range(len(dist)), p=dist)
            prediction = self.model.config.decoder.chars[idx]
            nodes += [prediction]
            if check_call:  # exception caught in main
                assert prediction not in ['DBranch', 'DExcept', 'DLoop', 'DSubTree']
            if prediction == 'STOP':
                edges += [SIBLING_EDGE]
                break
            js = self.generate_ast(psi, nodes, edges + [CHILD_EDGE])
            ast.append(js)
            edges += [SIBLING_EDGE]
            num += 1
        return ast, nodes, edges

    def generate_ast(self, psi, in_nodes=['DSubTree'], in_edges=[CHILD_EDGE]):
        ast = collections.OrderedDict()
        node = in_nodes[-1]

        # Return the "AST" if the node is an API call
        if node not in ['DBranch', 'DExcept', 'DLoop', 'DSubTree']:
            ast['node'] = 'DAPICall'
            ast['_call'] = node
            return ast

        ast['node'] = node
        nodes, edges = in_nodes[:], in_edges[:]

        if node == 'DBranch':
            ast_cond, nodes, edges = self.gen_until_STOP(psi, nodes, edges, check_call=True)
            ast_then, nodes, edges = self.gen_until_STOP(psi, nodes, edges)
            ast_else, nodes, edges = self.gen_until_STOP(psi, nodes, edges)
            ast['_cond'] = ast_cond
            ast['_then'] = ast_then
            ast['_else'] = ast_else
            return ast

        if node == 'DExcept':
            ast_try, nodes, edges = self.gen_until_STOP(psi, nodes, edges)
            ast_catch, nodes, edges = self.gen_until_STOP(psi, nodes, edges)
            ast['_try'] = ast_try
            ast['_catch'] = ast_catch
            return ast

        if node == 'DLoop':
            ast_cond, nodes, edges = self.gen_until_STOP(psi, nodes, edges, check_call=True)
            ast_body, nodes, edges = self.gen_until_STOP(psi, nodes, edges)
            ast['_cond'] = ast_cond
            ast['_body'] = ast_body
            return ast

        if node == 'DSubTree':
            ast_nodes, _, _ = self.gen_until_STOP(psi, nodes, edges)
            ast['_nodes'] = ast_nodes
            return ast

def find_api(nodes):
    for node in nodes:
        if node['node'] == 'DAPICall':
            call = node['_call'].split('.')
            api = '.'.join(call[:3])
            return api
    return None

def plot2d(asts):
    import matplotlib.pyplot as plt
    import matplotlib.cm as cm
    dic = {}
    for ast in asts:
        sample = ast['psi']
        api = find_api(ast['_nodes'])
        if api is None:
            continue
        if api not in dic:
            dic[api] = []
        dic[api].append(sample)

    apis = dic.keys()
    colors = cm.rainbow(np.linspace(0, 1, len(dic)))
    plotpoints = []
    for api, color in zip(apis, colors):
        x = list(map(lambda s: s[0], dic[api]))
        y = list(map(lambda s: s[1], dic[api]))
        plotpoints.append(plt.scatter(x, y, color=color))

    plt.legend(plotpoints, apis, scatterpoints=1, loc='lower left', ncol=3, fontsize=8)
    plt.axhline(0, color='black')
    plt.axvline(0, color='black')
    plt.show()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--save', type=str, default='save',
                        help='directory to laod model from')
    parser.add_argument('--evidence_file', type=str, default=None,
                        help='input file containing evidences (in JSON)')
    parser.add_argument('--random', action='store_true',
                        help='print random ASTs by sampling from Normal(0,1) (ignores evidences)')
    parser.add_argument('--plot2d', action='store_true',
                        help='plots the (2d) sampled psi values in scatterplot (requires --random)')
    parser.add_argument('--output_file', type=str, default=None,
                        help='file to print AST (in JSON) to')
    parser.add_argument('--n', type=int, default=1,
                        help='number of ASTs to sample/synthesize')

    clargs = parser.parse_args()
    if not clargs.evidence_file and not clargs.random:
        parser.error('Provide at least one option: --evidence_file or --random')
    if clargs.plot2d and not clargs.random:
        parser.error('--plot2d requires --random (otherwise there is only one psi to plot)')
    infer(clargs)
<commit_msg>Fixed bug in infer that caused long recursion<commit_after>from __future__ import print_function
import tensorflow as tf
import numpy as np

import argparse
import os
import json
import collections

from bayou.core.model import Model
from bayou.core.utils import CHILD_EDGE, SIBLING_EDGE
from bayou.core.utils import read_config

MAX_GEN_UNTIL_STOP = 20
MAX_AST_DEPTH = 5

def infer(clargs):
    with tf.Session() as sess:
        predictor = BayesianPredictor(clargs.save, sess)
        err = 0
        asts = []
        if clargs.evidence_file:
            with open(clargs.evidence_file) as f:
                js = json.load(f)
        for c in range(clargs.n):
            print('Generated {} ASTs ({} errors)'.format(c, err), end='\r')
            try:
                if clargs.random:
                    psi = predictor.psi_random()
                else:
                    psi = predictor.psi_from_evidence(js)
                ast = predictor.generate_ast(psi)
                if clargs.plot2d:
                    ast['psi'] = list(psi[0])
                asts.append(ast)
                c += 1
            except AssertionError:
                err += 1
    if clargs.plot2d:
        if predictor.model.config.latent_size == 2:
            plot2d(asts)
        else:
            print('Latent space is not 2-dimensional.. cannot plot')

    if clargs.output_file is None:
        print(json.dumps({'asts': asts}, indent=2))
    else:
        with open(clargs.output_file, 'w') as f:
            json.dump({'asts': asts}, fp=f, indent=2)
    print('Number of errors: {}'.format(err))

class BayesianPredictor(object):

    def __init__(self, save, sess):
        self.sess = sess

        # load the saved config
        with open(os.path.join(save, 'config.json')) as f:
            config = read_config(json.load(f), save_dir=save, infer=True)
        self.model = Model(config, True)

        # restore the saved model
        tf.global_variables_initializer().run()
        saver = tf.train.Saver(tf.global_variables())
        ckpt = tf.train.get_checkpoint_state(save)
        saver.restore(self.sess, ckpt.model_checkpoint_path)

    def infer(self, evidences):
        psi = self.psi_from_evidence(evidences)
        return self.generate_ast(psi)

    def psi_random(self):
        return np.random.normal(size=[1, self.model.config.latent_size])

    def psi_from_evidence(self, js_evidences):
        return self.model.infer_psi(self.sess, js_evidences)

    def gen_until_STOP(self, psi, depth, in_nodes, in_edges, check_call=False):
        ast = []
        nodes, edges = in_nodes[:], in_edges[:]
        num = 0
        while True:
            assert num < MAX_GEN_UNTIL_STOP # exception caught in main
            dist = self.model.infer_ast(self.sess, psi, nodes, edges)
            idx = np.random.choice(range(len(dist)), p=dist)
            prediction = self.model.config.decoder.chars[idx]
            nodes += [prediction]
            if check_call:  # exception caught in main
                assert prediction not in ['DBranch', 'DExcept', 'DLoop', 'DSubTree']
            if prediction == 'STOP':
                edges += [SIBLING_EDGE]
                break
            js = self.generate_ast(psi, depth + 1, nodes, edges + [CHILD_EDGE])
            ast.append(js)
            edges += [SIBLING_EDGE]
            num += 1
        return ast, nodes, edges

    def generate_ast(self, psi, depth=0, in_nodes=['DSubTree'], in_edges=[CHILD_EDGE]):
        assert depth < MAX_AST_DEPTH
        ast = collections.OrderedDict()
        node = in_nodes[-1]

        # Return the "AST" if the node is an API call
        if node not in ['DBranch', 'DExcept', 'DLoop', 'DSubTree']:
            ast['node'] = 'DAPICall'
            ast['_call'] = node
            return ast

        ast['node'] = node
        nodes, edges = in_nodes[:], in_edges[:]

        if node == 'DBranch':
            ast_cond, nodes, edges = self.gen_until_STOP(psi, depth, nodes, edges, check_call=True)
            ast_then, nodes, edges = self.gen_until_STOP(psi, depth, nodes, edges)
            ast_else, nodes, edges = self.gen_until_STOP(psi, depth, nodes, edges)
            ast['_cond'] = ast_cond
            ast['_then'] = ast_then
            ast['_else'] = ast_else
            return ast

        if node == 'DExcept':
            ast_try, nodes, edges = self.gen_until_STOP(psi, depth, nodes, edges)
            ast_catch, nodes, edges = self.gen_until_STOP(psi, depth, nodes, edges)
            ast['_try'] = ast_try
            ast['_catch'] = ast_catch
            return ast

        if node == 'DLoop':
            ast_cond, nodes, edges = self.gen_until_STOP(psi, depth, nodes, edges, check_call=True)
            ast_body, nodes, edges = self.gen_until_STOP(psi, depth, nodes, edges)
            ast['_cond'] = ast_cond
            ast['_body'] = ast_body
            return ast

        if node == 'DSubTree':
            ast_nodes, _, _ = self.gen_until_STOP(psi, depth, nodes, edges)
            ast['_nodes'] = ast_nodes
            return ast

def find_api(nodes):
    for node in nodes:
        if node['node'] == 'DAPICall':
            call = node['_call'].split('.')
            api = '.'.join(call[:3])
            return api
    return None

def plot2d(asts):
    import matplotlib.pyplot as plt
    import matplotlib.cm as cm
    dic = {}
    for ast in asts:
        sample = ast['psi']
        api = find_api(ast['_nodes'])
        if api is None:
            continue
        if api not in dic:
            dic[api] = []
        dic[api].append(sample)

    apis = dic.keys()
    colors = cm.rainbow(np.linspace(0, 1, len(dic)))
    plotpoints = []
    for api, color in zip(apis, colors):
        x = list(map(lambda s: s[0], dic[api]))
        y = list(map(lambda s: s[1], dic[api]))
        plotpoints.append(plt.scatter(x, y, color=color))

    plt.legend(plotpoints, apis, scatterpoints=1, loc='lower left', ncol=3, fontsize=8)
    plt.axhline(0, color='black')
    plt.axvline(0, color='black')
    plt.show()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--save', type=str, default='save',
                        help='directory to laod model from')
    parser.add_argument('--evidence_file', type=str, default=None,
                        help='input file containing evidences (in JSON)')
    parser.add_argument('--random', action='store_true',
                        help='print random ASTs by sampling from Normal(0,1) (ignores evidences)')
    parser.add_argument('--plot2d', action='store_true',
                        help='plots the (2d) sampled psi values in scatterplot (requires --random)')
    parser.add_argument('--output_file', type=str, default=None,
                        help='file to print AST (in JSON) to')
    parser.add_argument('--n', type=int, default=1,
                        help='number of ASTs to sample/synthesize')

    clargs = parser.parse_args()
    if not clargs.evidence_file and not clargs.random:
        parser.error('Provide at least one option: --evidence_file or --random')
    if clargs.plot2d and not clargs.random:
        parser.error('--plot2d requires --random (otherwise there is only one psi to plot)')
    infer(clargs)
<|endoftext|>


<commit_before>import numpy as np
from numpy import ma, random

from scipy.interpolate import RectBivariateSpline, interp1d

# ============================================================================
# I need to improve this, and include the places where the
#   flags are masked, i.e. only eliminate where the flags
#   could guarantee it was false.

def combined_flag(flags, criteria=None):
    """ Returns the combined flag considering all the criteria

        Input: flags

        Collects all flags in the criteria, and for each measurements, it
          return the maximum flag value among the different criteria.

        If criteria is not defined, considers all the flags,
          i.e. flags.keys()
    """
    assert hasattr(flags, 'keys')

    if criteria is None:
        criteria = list(flags.keys())

    output = np.asanyarray(flags[criteria[0]])
    for c in criteria[1:]:
        assert len(flags[c]) == len(output)
        output = np.max([output, flags[c]], axis=0)

    return output

def make_qc_index(flags, criteria, type="anytrue"):
    ind = flags[criteria[0]].copy()
    if type == "anytrue":
        for c in criteria:
            ind[(ind == True) | (flags[c] == True)] = True
        #ind[np.nonzero((ind == True) | (np.array(flags[c]) == True))[0]] = True
    elif type == "alltrue":
        for c in criteria:
            ind[(ind == True) | (flags[c] == True)] = True
    for c in criteria:
        ind[(ind == False) | (flags[c] == False)] = False
        #ind[np.nonzero((ind == False) | (np.array(flags[c]) == False))[0]] = False
    return ind
<commit_msg>Removing unecessary imports.<commit_after>import numpy as np

# ============================================================================
# I need to improve this, and include the places where the
#   flags are masked, i.e. only eliminate where the flags
#   could guarantee it was false.

def combined_flag(flags, criteria=None):
    """ Returns the combined flag considering all the criteria

        Input: flags

        Collects all flags in the criteria, and for each measurements, it
          return the maximum flag value among the different criteria.

        If criteria is not defined, considers all the flags,
          i.e. flags.keys()
    """
    assert hasattr(flags, 'keys')

    if criteria is None:
        criteria = list(flags.keys())

    output = np.asanyarray(flags[criteria[0]])
    for c in criteria[1:]:
        assert len(flags[c]) == len(output)
        output = np.max([output, flags[c]], axis=0)

    return output

def make_qc_index(flags, criteria, type="anytrue"):
    ind = flags[criteria[0]].copy()
    if type == "anytrue":
        for c in criteria:
            ind[(ind == True) | (flags[c] == True)] = True
        #ind[np.nonzero((ind == True) | (np.array(flags[c]) == True))[0]] = True
    elif type == "alltrue":
        for c in criteria:
            ind[(ind == True) | (flags[c] == True)] = True
    for c in criteria:
        ind[(ind == False) | (flags[c] == False)] = False
        #ind[np.nonzero((ind == False) | (np.array(flags[c]) == False))[0]] = False
    return ind
<|endoftext|>