<commit_before>#!/usr/bin/python
"""
  Light weight wrapper to Husky
  usage:
     ./husky.py <task> [<replay log file> [F|FF]]
"""

import sys
import serial
import struct
from logit import *
from crc16 import ccitt_checksum as checksum

PROTOCOL_VERSION = 1

REQ_PLATFORM_INFO = 0x4001
REQ_FIRMWARE_INFO = 0x4003
REQ_MAX_SPEED = 0x4210
REQ_MAX_ACCEL = 0x4211
REQ_GPIO = 0x4301
REQ_PLATFORM_ORIENTATION = 0x4600   # not supported?
REQ_PLATFORM_MAGNETOMETERS = 0x4606 # not supported?

SET_DIFFERENTIAL_OUTPUT = 0x0202

class Husky:
  def __init__( self, com ):
    self.com = com
    self.timestamp = 0
    self.cmdSpeed = (0, 0)
    self.enc = None
    self.config()

  def sendPacket( self, messageType, data = "" ):
    packet = struct.pack( '=BBBBIBHB',
    0xAA,           # SAH
    len(data)+11,   # data length + compliment
    255-len(data)-11,
    PROTOCOL_VERSION,  # version (in doc is used version 1, but sample code sends 0)
    self.timestamp, # unique number for ACK
    0,              # flags
    messageType,    # 16bit
    0x55 )          # STX - data start
    packet += data
    packet += struct.pack("H", checksum( [ord(x) for x in packet] ))
#    print repr(packet)
    self.timestamp += 1
    self.com.write( packet )

  def readPacket( self ):
    b = self.com.read(1)
    while b != chr(0xAA):
      print "skipping", hex(ord(b))
      b = self.com.read(1)
    length = ord(self.com.read(1))
    notLength = ord(self.com.read(1))
    assert length+notLength == 0xFF, (length, notLength)
    ret = ""
    for i in xrange(length):
      b = self.com.read(1)
      ret += b
    assert checksum( [0xAA, length, notLength] + [ord(x) for x in ret[:-2]] ) == ord(ret[-2]) + 256*ord(ret[-1])
    version, timestamp, flags, msgType, stx = struct.unpack("=BIBHB", ret[:9])
    assert version == 0x00, version # receiving ver0, but sending ver1
    assert flags == 0, flags
    assert stx == 0x55, stx
    return timestamp, msgType, ret[9:-2]

  def config( self ):
    self.sendPacket( REQ_PLATFORM_MAGNETOMETERS, data=struct.pack("h", 10 ) ) # at 10Hz
    self.sendPacket( REQ_FIRMWARE_INFO, data=struct.pack("h", 0 ) ) # once
    self.sendPacket( REQ_MAX_SPEED, data=struct.pack("h", 0 ) ) # once
    self.sendPacket( REQ_MAX_ACCEL, data=struct.pack("h", 0 ) ) # once
    self.sendPacket( REQ_GPIO, data=struct.pack("h", 0 ) ) # once
    self.sendPacket( REQ_PLATFORM_ORIENTATION, data=struct.pack("h", 10 ) ) # 10Hz

  def update( self ):
    self.sendPacket( SET_DIFFERENTIAL_OUTPUT, data=struct.pack("hh", self.cmdSpeed[0], self.cmdSpeed[1]) )
    for i in xrange(200):
      timestamp, msgType, data = self.readPacket()
      print hex(msgType)
      if msgType & 0xF000 == 0x4000: # confirmations
        print hex(msgType), [hex(ord(x)) for x in data]
      if msgType in [0x8210, 0x8211, 0x8003 ]:
        print hex(msgType), [hex(ord(x)) for x in data]
      if msgType == 0x8800: # Encoder data
        self.enc = struct.unpack( "=Biihh", data ) # expected 2 encoders - position and speed
#        print "ENC", self.enc
        break

def main0( com ):
  robot = Husky( com )
#  robot.sendPacket( REQ_PLATFORM_INFO )
#  robot.sendPacket( REQ_FIRMWARE_INFO, data=struct.pack("H",0) )
  robot.sendPacket( SET_DIFFERENTIAL_OUTPUT, data=struct.pack("hh", 3000, 3000) )
  for i in xrange(200):
    timestamp, msgType, data = robot.readPacket()
    if msgType == 0x8004:
      # System status data
      uptime, numMeas = struct.unpack( "=IB", data[:5] )
      data = data[5:]
      assert numMeas == 3, numMeas
      volt1, volt2, volt3, numMeas = struct.unpack( "hhhB", data[:6+1] )
      data = data[6+1:]
      print volt1/100., volt2/100., volt3/100., "|",
      assert numMeas == 3, numMeas
      cur1, cur2, cur3, numMeas = struct.unpack( "hhhB", data[:6+1] )
      data = data[6+1:]
      print cur1/100., cur2/100., cur3/100., "|",
      assert numMeas == 4, numMeas
      t1, t2, t3, t4 = struct.unpack( "hhhh", data )
      print t1/100., t2/100., t3/100., t4/100.
    if msgType == 0x8005:
      # Power system status data
      batNum, batState, batCap, batDesc = struct.unpack("=BHHB", data )
      assert batNum == 1, batNum
      assert batDesc == 0xC2, hex(batDesc) # present, in use, Lead-acid battery
      print batState, batCap
    if msgType == 0x8800:
      # Encoder data
      enc = struct.unpack( "=Biihh", data ) # expected 2 encoders - position and speed
      print "ENC", enc

def testMotion( com ):
  "slow motion for 25cm (after reseting encoders)"
  robot = Husky( com )
  robot.cmdSpeed = (1000, 1000)
  for i in xrange(1000):
    robot.update()
    if robot.enc and robot.enc[1] > 250:
      break
  robot.cmdSpeed = (0, 0)
  robot.update()

def testConfig( com ):
  "sensor configuration should be in Husky constructor - make sure you receive the data"
  robot = Husky( com )
  for i in xrange(10):
    robot.update()

if __name__ == "__main__":
  if len(sys.argv) < 2:
    print __doc__
    sys.exit(1)
  filename = None
  com = None
  if len(sys.argv) > 2:
    replayAssert = True
    filename = sys.argv[2]
    if len(sys.argv) > 3:
      assert sys.argv[3] in ['F','FF']
      if sys.argv[2] == 'F':
        replayAssert = False
      else:
        com = ReplyLogInputsOnly( filename )
  if filename:
    if com == None:
      com = ReplayLog( filename, assertWrite=replayAssert )
  else:
    com = LogIt( serial.Serial( '/dev/ttyUSB1', 115200 ) )
  testConfig( com )

<commit_msg>experimental test for going straight on Robotem Rovne 2015 (1m/s)<commit_after>#!/usr/bin/python
"""
  Light weight wrapper to Husky
  usage:
     ./husky.py <task> [<replay log file> [F|FF]]
"""

import sys
import serial
import struct
from logit import *
from crc16 import ccitt_checksum as checksum

PROTOCOL_VERSION = 1

REQ_PLATFORM_INFO = 0x4001
REQ_FIRMWARE_INFO = 0x4003
REQ_MAX_SPEED = 0x4210
REQ_MAX_ACCEL = 0x4211
REQ_GPIO = 0x4301
REQ_PLATFORM_ORIENTATION = 0x4600   # not supported?
REQ_PLATFORM_MAGNETOMETERS = 0x4606 # not supported?

SET_DIFFERENTIAL_OUTPUT = 0x0202

class Husky:
  def __init__( self, com ):
    self.com = com
    self.timestamp = 0
    self.cmdSpeed = (0, 0)
    self.enc = None
    self.config()

  def sendPacket( self, messageType, data = "" ):
    packet = struct.pack( '=BBBBIBHB',
    0xAA,           # SAH
    len(data)+11,   # data length + compliment
    255-len(data)-11,
    PROTOCOL_VERSION,  # version (in doc is used version 1, but sample code sends 0)
    self.timestamp, # unique number for ACK
    0,              # flags
    messageType,    # 16bit
    0x55 )          # STX - data start
    packet += data
    packet += struct.pack("H", checksum( [ord(x) for x in packet] ))
#    print repr(packet)
    self.timestamp += 1
    self.com.write( packet )

  def readPacket( self ):
    b = self.com.read(1)
    while b != chr(0xAA):
      print "skipping", hex(ord(b))
      b = self.com.read(1)
    length = ord(self.com.read(1))
    notLength = ord(self.com.read(1))
    assert length+notLength == 0xFF, (length, notLength)
    ret = ""
    for i in xrange(length):
      b = self.com.read(1)
      ret += b
    assert checksum( [0xAA, length, notLength] + [ord(x) for x in ret[:-2]] ) == ord(ret[-2]) + 256*ord(ret[-1])
    version, timestamp, flags, msgType, stx = struct.unpack("=BIBHB", ret[:9])
    assert version == 0x00, version # receiving ver0, but sending ver1
    assert flags == 0, flags
    assert stx == 0x55, stx
    return timestamp, msgType, ret[9:-2]

  def config( self ):
    self.sendPacket( REQ_PLATFORM_MAGNETOMETERS, data=struct.pack("h", 10 ) ) # at 10Hz
    self.sendPacket( REQ_FIRMWARE_INFO, data=struct.pack("h", 0 ) ) # once
    self.sendPacket( REQ_MAX_SPEED, data=struct.pack("h", 0 ) ) # once
    self.sendPacket( REQ_MAX_ACCEL, data=struct.pack("h", 0 ) ) # once
    self.sendPacket( REQ_GPIO, data=struct.pack("h", 0 ) ) # once
    self.sendPacket( REQ_PLATFORM_ORIENTATION, data=struct.pack("h", 10 ) ) # 10Hz

  def update( self ):
    self.sendPacket( SET_DIFFERENTIAL_OUTPUT, data=struct.pack("hh", self.cmdSpeed[0], self.cmdSpeed[1]) )
    for i in xrange(200):
      timestamp, msgType, data = self.readPacket()
      print hex(msgType)
      if msgType & 0xF000 == 0x4000: # confirmations
        print hex(msgType), [hex(ord(x)) for x in data]
      if msgType in [0x8210, 0x8211, 0x8003 ]:
        print hex(msgType), [hex(ord(x)) for x in data]
      if msgType == 0x8800: # Encoder data
        self.enc = struct.unpack( "=Biihh", data ) # expected 2 encoders - position and speed
#        print "ENC", self.enc
        break

def main0( com ):
  robot = Husky( com )
#  robot.sendPacket( REQ_PLATFORM_INFO )
#  robot.sendPacket( REQ_FIRMWARE_INFO, data=struct.pack("H",0) )
  robot.sendPacket( SET_DIFFERENTIAL_OUTPUT, data=struct.pack("hh", 3000, 3000) )
  for i in xrange(200):
    timestamp, msgType, data = robot.readPacket()
    if msgType == 0x8004:
      # System status data
      uptime, numMeas = struct.unpack( "=IB", data[:5] )
      data = data[5:]
      assert numMeas == 3, numMeas
      volt1, volt2, volt3, numMeas = struct.unpack( "hhhB", data[:6+1] )
      data = data[6+1:]
      print volt1/100., volt2/100., volt3/100., "|",
      assert numMeas == 3, numMeas
      cur1, cur2, cur3, numMeas = struct.unpack( "hhhB", data[:6+1] )
      data = data[6+1:]
      print cur1/100., cur2/100., cur3/100., "|",
      assert numMeas == 4, numMeas
      t1, t2, t3, t4 = struct.unpack( "hhhh", data )
      print t1/100., t2/100., t3/100., t4/100.
    if msgType == 0x8005:
      # Power system status data
      batNum, batState, batCap, batDesc = struct.unpack("=BHHB", data )
      assert batNum == 1, batNum
      assert batDesc == 0xC2, hex(batDesc) # present, in use, Lead-acid battery
      print batState, batCap
    if msgType == 0x8800:
      # Encoder data
      enc = struct.unpack( "=Biihh", data ) # expected 2 encoders - position and speed
      print "ENC", enc

def testMotion( com ):
  "slow motion for 25cm (after reseting encoders)"
  robot = Husky( com )
  robot.cmdSpeed = (1000, 1000)
  for i in xrange(1000):
    robot.update()
    if robot.enc and robot.enc[1] > 250:
      break
  robot.cmdSpeed = (0, 0)
  robot.update()

def testConfig( com ):
  "sensor configuration should be in Husky constructor - make sure you receive the data"
  robot = Husky( com )
  for i in xrange(10):
    robot.update()

def testRR2015( com ):
  "Robotem Rovne 2015"
  robot = Husky( com )
  robot.cmdSpeed = (10000, 10000)
  while True:
    robot.update()
  robot.cmdSpeed = (0, 0)
  robot.update()

if __name__ == "__main__":
  if len(sys.argv) < 2:
    print __doc__
    sys.exit(1)
  filename = None
  com = None
  if len(sys.argv) > 2:
    replayAssert = True
    filename = sys.argv[2]
    if len(sys.argv) > 3:
      assert sys.argv[3] in ['F','FF']
      if sys.argv[2] == 'F':
        replayAssert = False
      else:
        com = ReplyLogInputsOnly( filename )
  if filename:
    if com == None:
      com = ReplayLog( filename, assertWrite=replayAssert )
  else:
    com = LogIt( serial.Serial( '/dev/ttyUSB1', 115200 ) )
  testRR2015( com )

<|endoftext|>


<commit_before>import datetime
import sys
import time

import iso8601 as iso

import config
# This is required so that util.call_command works
import plugins, hooks  # pylint: disable=unused-import
import log
from utils import util
from utils import web

strptime = datetime.datetime.strptime

class Events(object):
    def __init__(self, bot):
        self.bot = bot
        self.config = self.bot.config
        self.userdb = self.bot.userdb

    def on_all(self, event, irc, arguments):
        if event.raw.find("%") == -1:
            log.debug(event.raw)

    def on_error(self, event, irc):
        log.error(" ".join(event.arguments))
        if " ".join(event.arguments).startswith("Closing Link"):
            self.web.stop()
            self.userdb.flush()
            self.db_job.stop()
            sys.exit(1)

    @staticmethod
    def on_ctcp(irc, event, raw):
        log.info("Received CTCP reply " + raw)

    def on_privmsg(self, event, irc, arguments):
        str_args = " ".join(arguments)
        if (str_args.startswith(config.commandChar) or
           str_args.startswith(self.config['nickname'])):
            util.call_command(self.bot, event, irc, arguments)
        else:
            util.call_hook(self.bot, event, irc, arguments)
        nick = event.source.nick
        str_args = ' '.join(arguments)
        self._update_seen_db(event, irc, nick, str_args)

    def _get_time(self, tags):
        if len(tags):
            for i in tags:
                try:
                    date = iso.parse_date(i['time'])
                    timestamp = time.mktime(date.timetuple())
                except KeyError:
                    pass
        else:
            timestamp = time.time()
        return timestamp

    def _update_seen_db(self, event, irc, nick, str_args):
        timestamp = self._get_time(event.tags)
        try:
            udb = self.userdb[event.target][nick]
            if udb['seen'] is None:
                udb['seen'] = []
            msg = utils.irc.strip_colors(str_args)
            udb["seen"].append({'time': timestamp, 'message': msg})
            compare = lambda m: m["time"]
            self.userdb[event.target][nick]['seen'] = sorted(udb["seen"],
                                                             key=compare,
                                                             reverse=True)[-5:]
        except KeyError:
            irc.send("WHO {0} nuhs%nhuac".format(event.target))

    @staticmethod
    def on_send(data):
        if data.find("%") == -1:
            log.debug(data)

    def on_nick(self, event, irc):
        nick = event.source.nick
        to_nick = event.arguments[0]
        for chan in self.userdb:
            chandb = self.userdb[chan]
            for u in chandb.values():
                if u['host'] == event.source.host:
                    self.userdb[chan][to_nick] = chandb[nick]
                    self.userdb[chan][to_nick]['hostmask'] = event.source
                    del self.userdb[chan][nick]
                    break
            break

    def on_quit(self, event, irc):
        nick = event.source.nick
        if nick == self.config['nickname']:
            web.app.stop()
            self.web.stop()
            self.userdb.flush()
            self.db_job.stop()
            sys.exit(1)
        else:
            for chan in self.userdb:
                for u in self.userdb[chan].values():
                    if u['host'] == event.source.host:
                        del self.userdb[chan][nick]
                        break
                break

    def on_kick(self, event, irc):
        nick = event.raw.split(" ")[3]
        if nick == self.config['nickname']:
            log.warning("Kicked from %s, trying to re-join", event.target)
            irc.join(event.target)
        else:
            self.userdb.remove_entry(event, nick)

    def on_part(self, event, irc):
        requested = "".join(event.arguments).startswith("requested")
        nick = event.source.nick
        if nick == self.config['nickname']:
            if requested:
                log.warning("Removed from %s, trying to re-join", event.target)
                irc.join(event.target)
            else:
                del self.userdb[event.target]
        else:
            self.userdb.remove_entry(event, nick)

    def on_join(self, event, irc):
        if event.source.nick == self.config['nickname']:
            log.info("Joining %s", event.target)
            if not event.target in self.userdb:
                self.userdb[event.target] = {}
            irc.send("WHO {0} nuhs%nhuac".format(event.target))
        else:
            irc.send("WHO {0} nuhs%nhuac".format(event.source.nick))

    @staticmethod
    def on_invite(event, irc):
        hostmask = event.source.host
        channel = event.arguments[0]
        if util.check_perms(hostmask, channel, trusted=True):
            log.info("Invited to %s by %s", channel, hostmask)
            irc.join(channel)

    def on_notice(self, event, irc):
        source = event.source.host
        if not event.target == "*":
            if not event.target == self.config['nickname']:
                channel = event.target
                log.info("Received channel notice from %s in %s",
                         source,
                         channel)
            else:
                log.info("Received private notice from %s", source)

    # Numeric events
    def on_unavailresource(self, event, irc):
        log.error("Nick unavailable, trying alternative")
        irc.nick(self.config['nickname'] + "_")
        self.config['nickname'] = self.config['nickname'] + "_"

    def on_nicknameinuse(self, event, irc):
        log.error("Nick already in use, trying alternative")
        irc.nick(self.config['nickname'] + "_")
        self.config['nickname'] = self.config['nickname'] + "_"

    @staticmethod
    def on_bannedfromchan(event, irc):
        s = event.raw.split(" ")
        channel = s[3]
        irc.notice("wolfy1339", "Banned from {0}".format(channel))
        log.warning("Banned from %s", channel)

    @staticmethod
    def on_endofmotd(event, irc):
        log.info("Received MOTD from network")

    @staticmethod
    def on_welcome(event, irc):
        log.info("Connected to network")

    def on_whoreply(self, event, irc, arguments):
        nick = arguments[4]
        if nick != "ChanServ":
            (ident, host) = arguments[1:3]
            channel = arguments[0]
            hostmask = "{0}!{1}@{2}".format(nick, ident, host)
            if "gateway" not in host or not util.is_ip_or_rdns(host):
                account = host.split("/")[-1].split('.')[-1]
            else:
                account = None
            self.userdb.add_entry(channel, nick, hostmask, account)

    def on_whospcrpl(self, event, irc, arguments):
        nick = arguments[3]
        if nick != "ChanServ":
            args = event.arguments
            (ident, host) = args[1:3]
            hostmask = "{0}!{1}@{2}".format(nick, ident, host)
            channel = args[0]
            account = args[4] if args[4] != "0" else None
            self.userdb.add_entry(channel, nick, hostmask, account)

    @staticmethod
    def on_315(event, irc):
        log.info("Received end of WHO reply from network")
<commit_msg>We don't need/use self here, maark this as a static method<commit_after>import datetime
import sys
import time

import iso8601 as iso

import config
# This is required so that util.call_command works
import plugins, hooks  # pylint: disable=unused-import
import log
from utils import util
from utils import web

strptime = datetime.datetime.strptime

class Events(object):
    def __init__(self, bot):
        self.bot = bot
        self.config = self.bot.config
        self.userdb = self.bot.userdb

    def on_all(self, event, irc, arguments):
        if event.raw.find("%") == -1:
            log.debug(event.raw)

    def on_error(self, event, irc):
        log.error(" ".join(event.arguments))
        if " ".join(event.arguments).startswith("Closing Link"):
            self.web.stop()
            self.userdb.flush()
            self.db_job.stop()
            sys.exit(1)

    @staticmethod
    def on_ctcp(irc, event, raw):
        log.info("Received CTCP reply " + raw)

    def on_privmsg(self, event, irc, arguments):
        str_args = " ".join(arguments)
        if (str_args.startswith(config.commandChar) or
           str_args.startswith(self.config['nickname'])):
            util.call_command(self.bot, event, irc, arguments)
        else:
            util.call_hook(self.bot, event, irc, arguments)
        nick = event.source.nick
        str_args = ' '.join(arguments)
        self._update_seen_db(event, irc, nick, str_args)

    @staticmethod
    def _get_time(tags):
        if len(tags):
            for i in tags:
                try:
                    date = iso.parse_date(i['time'])
                    timestamp = time.mktime(date.timetuple())
                except KeyError:
                    pass
        else:
            timestamp = time.time()
        return timestamp

    def _update_seen_db(self, event, irc, nick, str_args):
        timestamp = self._get_time(event.tags)
        try:
            udb = self.userdb[event.target][nick]
            if udb['seen'] is None:
                udb['seen'] = []
            msg = utils.irc.strip_colors(str_args)
            udb["seen"].append({'time': timestamp, 'message': msg})
            compare = lambda m: m["time"]
            self.userdb[event.target][nick]['seen'] = sorted(udb["seen"],
                                                             key=compare,
                                                             reverse=True)[-5:]
        except KeyError:
            irc.send("WHO {0} nuhs%nhuac".format(event.target))

    @staticmethod
    def on_send(data):
        if data.find("%") == -1:
            log.debug(data)

    def on_nick(self, event, irc):
        nick = event.source.nick
        to_nick = event.arguments[0]
        for chan in self.userdb:
            chandb = self.userdb[chan]
            for u in chandb.values():
                if u['host'] == event.source.host:
                    self.userdb[chan][to_nick] = chandb[nick]
                    self.userdb[chan][to_nick]['hostmask'] = event.source
                    del self.userdb[chan][nick]
                    break
            break

    def on_quit(self, event, irc):
        nick = event.source.nick
        if nick == self.config['nickname']:
            web.app.stop()
            self.web.stop()
            self.userdb.flush()
            self.db_job.stop()
            sys.exit(1)
        else:
            for chan in self.userdb:
                for u in self.userdb[chan].values():
                    if u['host'] == event.source.host:
                        del self.userdb[chan][nick]
                        break
                break

    def on_kick(self, event, irc):
        nick = event.raw.split(" ")[3]
        if nick == self.config['nickname']:
            log.warning("Kicked from %s, trying to re-join", event.target)
            irc.join(event.target)
        else:
            self.userdb.remove_entry(event, nick)

    def on_part(self, event, irc):
        requested = "".join(event.arguments).startswith("requested")
        nick = event.source.nick
        if nick == self.config['nickname']:
            if requested:
                log.warning("Removed from %s, trying to re-join", event.target)
                irc.join(event.target)
            else:
                del self.userdb[event.target]
        else:
            self.userdb.remove_entry(event, nick)

    def on_join(self, event, irc):
        if event.source.nick == self.config['nickname']:
            log.info("Joining %s", event.target)
            if not event.target in self.userdb:
                self.userdb[event.target] = {}
            irc.send("WHO {0} nuhs%nhuac".format(event.target))
        else:
            irc.send("WHO {0} nuhs%nhuac".format(event.source.nick))

    @staticmethod
    def on_invite(event, irc):
        hostmask = event.source.host
        channel = event.arguments[0]
        if util.check_perms(hostmask, channel, trusted=True):
            log.info("Invited to %s by %s", channel, hostmask)
            irc.join(channel)

    def on_notice(self, event, irc):
        source = event.source.host
        if not event.target == "*":
            if not event.target == self.config['nickname']:
                channel = event.target
                log.info("Received channel notice from %s in %s",
                         source,
                         channel)
            else:
                log.info("Received private notice from %s", source)

    # Numeric events
    def on_unavailresource(self, event, irc):
        log.error("Nick unavailable, trying alternative")
        irc.nick(self.config['nickname'] + "_")
        self.config['nickname'] = self.config['nickname'] + "_"

    def on_nicknameinuse(self, event, irc):
        log.error("Nick already in use, trying alternative")
        irc.nick(self.config['nickname'] + "_")
        self.config['nickname'] = self.config['nickname'] + "_"

    @staticmethod
    def on_bannedfromchan(event, irc):
        s = event.raw.split(" ")
        channel = s[3]
        irc.notice("wolfy1339", "Banned from {0}".format(channel))
        log.warning("Banned from %s", channel)

    @staticmethod
    def on_endofmotd(event, irc):
        log.info("Received MOTD from network")

    @staticmethod
    def on_welcome(event, irc):
        log.info("Connected to network")

    def on_whoreply(self, event, irc, arguments):
        nick = arguments[4]
        if nick != "ChanServ":
            (ident, host) = arguments[1:3]
            channel = arguments[0]
            hostmask = "{0}!{1}@{2}".format(nick, ident, host)
            if "gateway" not in host or not util.is_ip_or_rdns(host):
                account = host.split("/")[-1].split('.')[-1]
            else:
                account = None
            self.userdb.add_entry(channel, nick, hostmask, account)

    def on_whospcrpl(self, event, irc, arguments):
        nick = arguments[3]
        if nick != "ChanServ":
            args = event.arguments
            (ident, host) = args[1:3]
            hostmask = "{0}!{1}@{2}".format(nick, ident, host)
            channel = args[0]
            account = args[4] if args[4] != "0" else None
            self.userdb.add_entry(channel, nick, hostmask, account)

    @staticmethod
    def on_315(event, irc):
        log.info("Received end of WHO reply from network")
<|endoftext|>


<commit_before>import datetime
import sys
import time

import iso8601 as iso

import config
# This is required so that util.call_command works
import plugins, hooks  # pylint: disable=unused-import
import log
from utils import util

strptime = datetime.datetime.strptime

class Events(object):
    def __init__(self, bot):
        self.bot = bot
        self.config = self.bot.config
        self.userdb = self.bot.userdb

    def on_all(self, event, irc, arguments):
        if event.raw.find("%") == -1:
            log.debug(event.raw)

    def on_error(self, event, irc):
        log.error(" ".join(event.arguments))
        if " ".join(event.arguments).startswith("Closing Link"):
            self.web.stop()
            self.userdb.flush()
            self.db_job.stop()
            sys.exit(1)

    @staticmethod
    def on_ctcp(irc, event, raw):
        log.info("Received CTCP reply " + raw)

    def on_privmsg(self, event, irc, arguments):
        str_args = " ".join(arguments)
        if (str_args.startswith(config.commandChar) or
           str_args.startswith(self.config['nickname'])):
            util.call_command(self.bot, event, irc, arguments)
        else:
            util.call_hook(self.bot, event, irc, arguments)
        nick = event.source.nick
        str_args = ' '.join(arguments)
        self._update_seen_db(event, nick, str_args)

    def _get_time(self, tags):
        if len(tags):
            for i in tags:
                try:
                    date = iso.parse_date(i['time'])
                    timestamp = time.mktime(date.timetuple())
                except KeyError:
                    pass
        else:
            timestamp = time.time()

    def _update_seen_db(self, event, nick, str_args):
        timestamp = self._get_time(event.tags)
        try:
            udb = self.userdb[event.target][nick]
            udb["seen"].append({'time': timestamp, 'messsage': str_args})
            compare = lambda m: m["time"]
            self.userdb[event.target][nick]['seen'] = sorted(udb["seen"],
                                                             key=compare,
                                                             reverse=True)[:-5]
        except KeyError:
            irc.send("WHO {0} nuhs%nhuac".format(event.target))

    @staticmethod
    def on_send(data):
        if data.find("%") == -1:
            log.debug(data)

    def on_nick(self, event, irc):
        nick = event.source.nick
        to_nick = event.arguments[0]
        for chan in self.userdb:
            chandb = self.userdb[chan]
            for u in chandb.values():
                if u['host'] == event.source.host:
                    self.userdb[chan][to_nick] = chandb[nick]
                    self.userdb[chan][to_nick]['hostmask'] = event.source
                    del self.userdb[chan][nick]
                    break
            break

    def on_quit(self, event, irc):
        nick = event.source.nick
        if nick == self.config['nickname']:
            self.web.stop()
            self.userdb.flush()
            self.db_job.stop()
            sys.exit(1)
        else:
            for chan in self.userdb:
                for u in self.userdb[chan].values():
                    if u['host'] == event.source.host:
                        del self.userdb[chan][nick]
                        break
                break

    def on_kick(self, event, irc):
        nick = event.raw.split(" ")[3]
        if nick == self.config['nickname']:
            log.warning("Kicked from %s, trying to re-join", event.target)
            irc.join(event.target)
        else:
            self.userdb.remove_entry(event, nick)

    def on_part(self, event, irc):
        requested = "".join(event.arguments).startswith("requested")
        nick = event.source.nick
        if nick == self.config['nickname']:
            if requested:
                log.warning("Removed from %s, trying to re-join", event.target)
                irc.join(event.target)
            else:
                del self.userdb[event.target]
        else:
            self.userdb.remove_entry(event, nick)

    def on_join(self, event, irc):
        if event.source.nick == self.config['nickname']:
            log.info("Joining %s", event.target)
            if not event.target in self.userdb:
                self.userdb[event.target] = {}
            irc.send("WHO {0} nuhs%nhuac".format(event.target))
        else:
            irc.send("WHO {0} nuhs%nhuac".format(event.source.nick))

    @staticmethod
    def on_invite(event, irc):
        hostmask = event.source.host
        channel = event.arguments[0]
        if util.check_perms(hostmask, channel, trusted=True):
            log.info("Invited to %s by %s", channel, hostmask)
            irc.join(channel)

    def on_notice(self, event, irc):
        source = event.source.host
        if not event.target == "*":
            if not event.target == self.config['nickname']:
                channel = event.target
                log.info("Received channel notice from %s in %s", source, channel)
            else:
                log.info("Received private notice from %s", source)

    # Numeric events
    def on_unavailresource(self, event, irc):
        log.error("Nick unavailable, trying alternative")
        irc.nick(self.config['nickname'] + "_")
        self.config['nickname'] = self.config['nickname'] + "_"

    def on_nicknameinuse(self, event, irc):
        log.error("Nick already in use, trying alternative")
        irc.nick(self.config['nickname'] + "_")
        self.config['nickname'] = self.config['nickname'] + "_"

    @staticmethod
    def on_bannedfromchan(event, irc):
        s = event.raw.split(" ")
        channel = s[3]
        irc.notice("wolfy1339", "Banned from {0}".format(channel))
        log.warning("Banned from %s", channel)

    @staticmethod
    def on_endofmotd(event, irc):
        log.info("Received MOTD from network")

    @staticmethod
    def on_welcome(event, irc):
        log.info("Connected to network")

    def on_whoreply(self, event, irc, arguments):
        nick = arguments[4]
        if nick != "ChanServ":
            (ident, host) = arguments[1:3]
            channel = arguments[0]
            hostmask = "{0}!{1}@{2}".format(nick, ident, host)
            if "gateway" not in host or not util.is_ip_or_rdns(host):
                account = host.split("/")[-1].split('.')[-1]
            else:
                account = None
            self.userdb.add_entry(channel, nick, hostmask, account)

    def on_whospcrpl(self, event, irc, arguments):
        nick = arguments[3]
        if nick != "ChanServ":
            args = event.arguments
            (ident, host) = args[1:3]
            hostmask = "{0}!{1}@{2}".format(nick, ident, host)
            channel = args[0]
            account = args[4] if args[4] != "0" else None
            self.userdb.add_entry(channel, nick, hostmask, account)

    @staticmethod
    def on_315(event, irc):
        log.info("Received end of WHO reply from network")
<commit_msg>Fix some issues<commit_after>import datetime
import sys
import time

import iso8601 as iso

import config
# This is required so that util.call_command works
import plugins, hooks  # pylint: disable=unused-import
import log
from utils import util

strptime = datetime.datetime.strptime

class Events(object):
    def __init__(self, bot):
        self.bot = bot
        self.config = self.bot.config
        self.userdb = self.bot.userdb

    def on_all(self, event, irc, arguments):
        if event.raw.find("%") == -1:
            log.debug(event.raw)

    def on_error(self, event, irc):
        log.error(" ".join(event.arguments))
        if " ".join(event.arguments).startswith("Closing Link"):
            self.web.stop()
            self.userdb.flush()
            self.db_job.stop()
            sys.exit(1)

    @staticmethod
    def on_ctcp(irc, event, raw):
        log.info("Received CTCP reply " + raw)

    def on_privmsg(self, event, irc, arguments):
        str_args = " ".join(arguments)
        if (str_args.startswith(config.commandChar) or
           str_args.startswith(self.config['nickname'])):
            util.call_command(self.bot, event, irc, arguments)
        else:
            util.call_hook(self.bot, event, irc, arguments)
        nick = event.source.nick
        str_args = ' '.join(arguments)
        self._update_seen_db(event, irc, nick, str_args)

    def _get_time(self, tags):
        if len(tags):
            for i in tags:
                try:
                    date = iso.parse_date(i['time'])
                    timestamp = time.mktime(date.timetuple())
                except KeyError:
                    pass
        else:
            timestamp = time.time()
        return timestamp

    def _update_seen_db(self, event, irc, nick, str_args):
        timestamp = self._get_time(event.tags)
        try:
            udb = self.userdb[event.target][nick]
            udb["seen"].append({'time': timestamp, 'messsage': str_args})
            compare = lambda m: m["time"]
            self.userdb[event.target][nick]['seen'] = sorted(udb["seen"],
                                                             key=compare,
                                                             reverse=True)[:-5]
        except KeyError:
            irc.send("WHO {0} nuhs%nhuac".format(event.target))

    @staticmethod
    def on_send(data):
        if data.find("%") == -1:
            log.debug(data)

    def on_nick(self, event, irc):
        nick = event.source.nick
        to_nick = event.arguments[0]
        for chan in self.userdb:
            chandb = self.userdb[chan]
            for u in chandb.values():
                if u['host'] == event.source.host:
                    self.userdb[chan][to_nick] = chandb[nick]
                    self.userdb[chan][to_nick]['hostmask'] = event.source
                    del self.userdb[chan][nick]
                    break
            break

    def on_quit(self, event, irc):
        nick = event.source.nick
        if nick == self.config['nickname']:
            self.web.stop()
            self.userdb.flush()
            self.db_job.stop()
            sys.exit(1)
        else:
            for chan in self.userdb:
                for u in self.userdb[chan].values():
                    if u['host'] == event.source.host:
                        del self.userdb[chan][nick]
                        break
                break

    def on_kick(self, event, irc):
        nick = event.raw.split(" ")[3]
        if nick == self.config['nickname']:
            log.warning("Kicked from %s, trying to re-join", event.target)
            irc.join(event.target)
        else:
            self.userdb.remove_entry(event, nick)

    def on_part(self, event, irc):
        requested = "".join(event.arguments).startswith("requested")
        nick = event.source.nick
        if nick == self.config['nickname']:
            if requested:
                log.warning("Removed from %s, trying to re-join", event.target)
                irc.join(event.target)
            else:
                del self.userdb[event.target]
        else:
            self.userdb.remove_entry(event, nick)

    def on_join(self, event, irc):
        if event.source.nick == self.config['nickname']:
            log.info("Joining %s", event.target)
            if not event.target in self.userdb:
                self.userdb[event.target] = {}
            irc.send("WHO {0} nuhs%nhuac".format(event.target))
        else:
            irc.send("WHO {0} nuhs%nhuac".format(event.source.nick))

    @staticmethod
    def on_invite(event, irc):
        hostmask = event.source.host
        channel = event.arguments[0]
        if util.check_perms(hostmask, channel, trusted=True):
            log.info("Invited to %s by %s", channel, hostmask)
            irc.join(channel)

    def on_notice(self, event, irc):
        source = event.source.host
        if not event.target == "*":
            if not event.target == self.config['nickname']:
                channel = event.target
                log.info("Received channel notice from %s in %s", source, channel)
            else:
                log.info("Received private notice from %s", source)

    # Numeric events
    def on_unavailresource(self, event, irc):
        log.error("Nick unavailable, trying alternative")
        irc.nick(self.config['nickname'] + "_")
        self.config['nickname'] = self.config['nickname'] + "_"

    def on_nicknameinuse(self, event, irc):
        log.error("Nick already in use, trying alternative")
        irc.nick(self.config['nickname'] + "_")
        self.config['nickname'] = self.config['nickname'] + "_"

    @staticmethod
    def on_bannedfromchan(event, irc):
        s = event.raw.split(" ")
        channel = s[3]
        irc.notice("wolfy1339", "Banned from {0}".format(channel))
        log.warning("Banned from %s", channel)

    @staticmethod
    def on_endofmotd(event, irc):
        log.info("Received MOTD from network")

    @staticmethod
    def on_welcome(event, irc):
        log.info("Connected to network")

    def on_whoreply(self, event, irc, arguments):
        nick = arguments[4]
        if nick != "ChanServ":
            (ident, host) = arguments[1:3]
            channel = arguments[0]
            hostmask = "{0}!{1}@{2}".format(nick, ident, host)
            if "gateway" not in host or not util.is_ip_or_rdns(host):
                account = host.split("/")[-1].split('.')[-1]
            else:
                account = None
            self.userdb.add_entry(channel, nick, hostmask, account)

    def on_whospcrpl(self, event, irc, arguments):
        nick = arguments[3]
        if nick != "ChanServ":
            args = event.arguments
            (ident, host) = args[1:3]
            hostmask = "{0}!{1}@{2}".format(nick, ident, host)
            channel = args[0]
            account = args[4] if args[4] != "0" else None
            self.userdb.add_entry(channel, nick, hostmask, account)

    @staticmethod
    def on_315(event, irc):
        log.info("Received end of WHO reply from network")
<|endoftext|>


<commit_before>from constants import LIBRARY_FOLDER, THUMBNAILS, IMAGE_TYPE, PDF_TYPE, EPUB_TYPE, COMIC_TYPE
from datetime import datetime
import dateutil.parser
import epub
import imghdr
from models import Comic, ComicPage, File, Epub, Pdf
import os
from PIL import Image
from PyPDF2 import PdfFileReader
from rarfile import RarFile
import re
from shutil import copyfile
from tarfile import TarFile
from utils import add_and_refresh, path_leaf
from zipfile import ZipFile

class BaseHandler:
    def __init__(self):
        pass

    @staticmethod
    def _make_cover_thumbnail(cover_name, cover_path, book_name):
        cover_extension = os.path.splitext(cover_name)[1]
        thumbnail_name = '{0}{1}'.format(book_name, cover_extension)
        thumbnail_path = os.path.join(THUMBNAILS, thumbnail_name)
        if os.path.exists(thumbnail_path):
            os.remove(thumbnail_path)
        im = Image.open(cover_path)
        im.thumbnail((200, 200))
        im.save(thumbnail_path)
        thumbnail = File(name=cover_name, path=thumbnail_path, size=os.path.getsize(thumbnail_path), type=IMAGE_TYPE)
        add_and_refresh(thumbnail)
        return thumbnail

    @staticmethod
    def _save_original_file(uploaded_file, file_type):
        file_name = path_leaf(uploaded_file)
        library_path = os.path.join(LIBRARY_FOLDER, file_name)
        copyfile(uploaded_file, library_path)
        size = os.path.getsize(library_path)
        new_file = File(name=file_name, path=library_path, size=size, type=file_type)
        add_and_refresh(new_file)
        return new_file

class ComicHandler(BaseHandler):
    def _get_info_filename(self, info):
        return info.filename

    def _get_info_size(self, info):
        return info.file_size

    def _get_info_list(self, archive):
        return archive.infolist()

    def handle_file(self, name, comic_file_path):
        comic_file = self._save_original_file(comic_file_path, 'comic')
        comic = Comic(name=name, type=COMIC_TYPE, file_id=comic_file.id)
        match = re.match("(.+?)\s+(\d+)\.", comic_file_path)
        if match:
            comic.series = match.group(1)
            comic.issue = match.group(2)
        add_and_refresh(comic)
        extracted_path = os.path.join(LIBRARY_FOLDER, name)
        self.extract_file(comic, comic_file_path, extracted_path)

    def _extract_archive_info(self, archive, comic, extracted_path):
        page_number = 1
        info_list = self._get_info_list(archive)
        is_zip = isinstance(archive, ZipFile)
        for info in info_list:
            info_file_name = self._get_info_filename(info)
            info_head = os.path.split(info_file_name)[0]
            if info_head:
                extracted_sub_dir = os.path.join(extracted_path, info_head)
                if not os.path.isdir(extracted_sub_dir):
                    os.makedirs(extracted_sub_dir)
            page_path = os.path.join(extracted_path, info_file_name)
            if not is_zip and info.isdir():
                continue
            archive.extract(info_file_name, extracted_path)
            page_type = imghdr.what(page_path)
            if page_type:
                page_file_name = path_leaf(info_file_name)
                page_path = os.path.join(extracted_path, info_file_name)
                page_size = self._get_info_size(info)
                page_file = File(name=page_file_name, path=page_path, size=page_size, type=IMAGE_TYPE)
                add_and_refresh(page_file)
                page = ComicPage(page_number=page_number, comic_id=comic.id, file_id=page_file.id)
                page_number += 1
                add_and_refresh(page)
        first_page_index = 0
        first_page = self._get_info_filename(info_list[first_page_index])
        while not is_zip and info_list[first_page_index].isdir():
            first_page_index += 1
            first_page = self._get_info_filename(info_list[first_page_index])
        first_page_path = os.path.join(extracted_path, first_page)
        cover = self._make_cover_thumbnail(first_page, first_page_path, comic.name)
        comic.cover_id = cover.id

class CbzHandler(ComicHandler):
    def extract_file(self, comic, comic_file_path, extracted_path):
        with ZipFile(comic_file_path) as zip_file:
            self._extract_archive_info(zip_file, comic, extracted_path)

class CbrHandler(ComicHandler):
    def extract_file(self, comic, comic_file_path, extracted_path):
        with RarFile(comic_file_path) as rar_file:
            self._extract_archive_info(rar_file, comic, extracted_path)

class CbtHandler(ComicHandler):
    def _get_info_filename(self, info):
        return info.name

    def _get_info_size(self, info):
        return info.size

    def _get_info_list(self, archive):
        return archive.getmembers()

    def extract_file(self, comic, comic_file_path, extracted_path):
        with TarFile(comic_file_path) as tar_file:
            self._extract_archive_info(tar_file, comic, extracted_path)

class EpubHandler(BaseHandler):
    def handle_file(self, name, epub_file_path):
        epub_file = self._save_original_file(epub_file_path, 'epub')
        extracted_path = os.path.join(LIBRARY_FOLDER, name)
        self._extract_zip(epub_file.path, extracted_path)
        (title, creators, publisher, release_date, cover_id) = self._read_epub_meta(epub_file.path, extracted_path)
        epub_obj = Epub(name=title, file_id=epub_file.id, type=EPUB_TYPE, extracted_path=extracted_path,
                        author=creators, publisher=publisher, release_date=release_date, cover_id=cover_id)
        add_and_refresh(epub_obj)

    @staticmethod
    def _read_epub_meta(epub_library_file, extracted_path):
        with epub.open_epub(epub_library_file, 'r') as epub_file:
            metadata = epub_file.opf.metadata
            title = metadata.titles[0][0]                                       # (title, lang)
            creators = ', '.join([x[0] for x in metadata.creators])             # (name, role, file_as)
            publisher = metadata.publisher
            release_date = None
            if metadata.dates:
                try:
                    release_date = dateutil.parser.parse(metadata.dates[0][0])  # (date, event)
                except (ValueError, OverflowError):
                    pass

            cover_id = next((x[1] for x in metadata.metas if x[0] == 'cover'), None)     # {name, value)
            thumbnail_id = None
            if cover_id:
                cover_manifest = epub_file.get_item(cover_id)
                cover_rel_path = cover_manifest.href
                if os.name == 'nt':
                    cover_rel_path = cover_rel_path.replace("/", "\\")
                cover_file_name = path_leaf(cover_rel_path)
                cover_path = os.path.join(extracted_path, epub_file.content_path, cover_rel_path)
                thumbnail = BaseHandler._make_cover_thumbnail(cover_file_name, cover_path, title)
                thumbnail_id = thumbnail.id
            return title, creators, publisher, release_date, thumbnail_id

    @staticmethod
    def _extract_zip(epub_library_file, extracted_path):
        with ZipFile(epub_library_file) as archive:
            archive.extractall(extracted_path)

class PdfHandler(BaseHandler):
    def handle_file(self, name, pdf_file_path):
        pdf_file = self._save_original_file(pdf_file_path, 'pdf')
        (title, author, release_date) = self._read_pdf_meta(pdf_file.path, name)
        pdf = Pdf(name=title, file_id=pdf_file.id, type=PDF_TYPE, author=author, release_date=release_date)
        add_and_refresh(pdf)

    @staticmethod
    def _read_pdf_meta(pdf_file_path, name):
        pdf_stream = open(pdf_file_path, "rb")
        pdf_reader = PdfFileReader(pdf_stream)
        pdf_info = pdf_reader.getDocumentInfo()
        title = name
        author = None
        release_date = None
        if '/Title' in pdf_info:
            title = pdf_info['/Title']
        if '/Author' in pdf_info:
            author = pdf_info['/Author']
        if '/CreationDate' in pdf_info:
            date_string = pdf_info['/CreationDate'].replace("'", '')        # D:YYYYMMDDHHmmSSOHH'mm'
            release_date = datetime.strptime(date_string, 'D:%Y%m%d%H%M%S')
        return title, author, release_date

<commit_msg>Now moves uploaded file rather than copy<commit_after>from constants import LIBRARY_FOLDER, THUMBNAILS, IMAGE_TYPE, PDF_TYPE, EPUB_TYPE, COMIC_TYPE
from datetime import datetime
import dateutil.parser
import epub
import imghdr
from models import Comic, ComicPage, File, Epub, Pdf
import os
from PIL import Image
from PyPDF2 import PdfFileReader
from rarfile import RarFile
import re
from tarfile import TarFile
from utils import add_and_refresh, path_leaf
from zipfile import ZipFile

class BaseHandler:
    def __init__(self):
        pass

    @staticmethod
    def _make_cover_thumbnail(cover_name, cover_path, book_name):
        cover_extension = os.path.splitext(cover_name)[1]
        thumbnail_name = '{0}{1}'.format(book_name, cover_extension)
        thumbnail_path = os.path.join(THUMBNAILS, thumbnail_name)
        if os.path.exists(thumbnail_path):
            os.remove(thumbnail_path)
        im = Image.open(cover_path)
        im.thumbnail((200, 200))
        im.save(thumbnail_path)
        thumbnail = File(name=cover_name, path=thumbnail_path, size=os.path.getsize(thumbnail_path), type=IMAGE_TYPE)
        add_and_refresh(thumbnail)
        return thumbnail

    @staticmethod
    def _save_original_file(uploaded_file, file_type):
        file_name = path_leaf(uploaded_file)
        library_path = os.path.join(LIBRARY_FOLDER, file_name)
        os.rename(uploaded_file, library_path)
        size = os.path.getsize(library_path)
        new_file = File(name=file_name, path=library_path, size=size, type=file_type)
        add_and_refresh(new_file)
        return new_file

class ComicHandler(BaseHandler):
    def _get_info_filename(self, info):
        return info.filename

    def _get_info_size(self, info):
        return info.file_size

    def _get_info_list(self, archive):
        return archive.infolist()

    def handle_file(self, name, upload_file_path):
        comic_file = self._save_original_file(upload_file_path, COMIC_TYPE)
        comic = Comic(name=name, type=COMIC_TYPE, file_id=comic_file.id)
        match = re.match("(.+?)\s+(\d+)\.", upload_file_path)
        if match:
            comic.series = match.group(1)
            comic.issue = match.group(2)
        add_and_refresh(comic)
        extracted_path = os.path.join(LIBRARY_FOLDER, name)
        self.extract_file(comic, comic_file.path, extracted_path)

    def _extract_archive_info(self, archive, comic, extracted_path):
        page_number = 1
        info_list = self._get_info_list(archive)
        is_zip = isinstance(archive, ZipFile)
        for info in info_list:
            info_file_name = self._get_info_filename(info)
            info_head = os.path.split(info_file_name)[0]
            if info_head:
                extracted_sub_dir = os.path.join(extracted_path, info_head)
                if not os.path.isdir(extracted_sub_dir):
                    os.makedirs(extracted_sub_dir)
            page_path = os.path.join(extracted_path, info_file_name)
            if not is_zip and info.isdir():
                continue
            archive.extract(info_file_name, extracted_path)
            page_type = imghdr.what(page_path)
            if page_type:
                page_file_name = path_leaf(info_file_name)
                page_path = os.path.join(extracted_path, info_file_name)
                page_size = self._get_info_size(info)
                page_file = File(name=page_file_name, path=page_path, size=page_size, type=IMAGE_TYPE)
                add_and_refresh(page_file)
                page = ComicPage(page_number=page_number, comic_id=comic.id, file_id=page_file.id)
                page_number += 1
                add_and_refresh(page)
        first_page_index = 0
        first_page = self._get_info_filename(info_list[first_page_index])
        while not is_zip and info_list[first_page_index].isdir():
            first_page_index += 1
            first_page = self._get_info_filename(info_list[first_page_index])
        first_page_path = os.path.join(extracted_path, first_page)
        cover = self._make_cover_thumbnail(first_page, first_page_path, comic.name)
        comic.cover_id = cover.id

class CbzHandler(ComicHandler):
    def extract_file(self, comic, comic_file_path, extracted_path):
        with ZipFile(comic_file_path) as zip_file:
            self._extract_archive_info(zip_file, comic, extracted_path)

class CbrHandler(ComicHandler):
    def extract_file(self, comic, comic_file_path, extracted_path):
        with RarFile(comic_file_path) as rar_file:
            self._extract_archive_info(rar_file, comic, extracted_path)

class CbtHandler(ComicHandler):
    def _get_info_filename(self, info):
        return info.name

    def _get_info_size(self, info):
        return info.size

    def _get_info_list(self, archive):
        return archive.getmembers()

    def extract_file(self, comic, comic_file_path, extracted_path):
        with TarFile(comic_file_path) as tar_file:
            self._extract_archive_info(tar_file, comic, extracted_path)

class EpubHandler(BaseHandler):
    def handle_file(self, name, epub_file_path):
        epub_file = self._save_original_file(epub_file_path, EPUB_TYPE)
        extracted_path = os.path.join(LIBRARY_FOLDER, name)
        self._extract_zip(epub_file.path, extracted_path)
        (title, creators, publisher, release_date, cover_id) = self._read_epub_meta(epub_file.path, extracted_path)
        epub_obj = Epub(name=title, file_id=epub_file.id, type=EPUB_TYPE, extracted_path=extracted_path,
                        author=creators, publisher=publisher, release_date=release_date, cover_id=cover_id)
        add_and_refresh(epub_obj)

    @staticmethod
    def _read_epub_meta(epub_library_file, extracted_path):
        with epub.open_epub(epub_library_file, 'r') as epub_file:
            metadata = epub_file.opf.metadata
            title = metadata.titles[0][0]                                       # (title, lang)
            creators = ', '.join([x[0] for x in metadata.creators])             # (name, role, file_as)
            publisher = metadata.publisher
            release_date = None
            if metadata.dates:
                try:
                    release_date = dateutil.parser.parse(metadata.dates[0][0])  # (date, event)
                except (ValueError, OverflowError):
                    pass

            cover_id = next((x[1] for x in metadata.metas if x[0] == 'cover'), None)     # {name, value)
            thumbnail_id = None
            if cover_id:
                cover_manifest = epub_file.get_item(cover_id)
                cover_rel_path = cover_manifest.href
                if os.name == 'nt':
                    cover_rel_path = cover_rel_path.replace("/", "\\")
                cover_file_name = path_leaf(cover_rel_path)
                cover_path = os.path.join(extracted_path, epub_file.content_path, cover_rel_path)
                thumbnail = BaseHandler._make_cover_thumbnail(cover_file_name, cover_path, title)
                thumbnail_id = thumbnail.id
            return title, creators, publisher, release_date, thumbnail_id

    @staticmethod
    def _extract_zip(epub_library_file, extracted_path):
        with ZipFile(epub_library_file) as archive:
            archive.extractall(extracted_path)

class PdfHandler(BaseHandler):
    def handle_file(self, name, pdf_file_path):
        pdf_file = self._save_original_file(pdf_file_path, PDF_TYPE)
        (title, author, release_date) = self._read_pdf_meta(pdf_file.path, name)
        pdf = Pdf(name=title, file_id=pdf_file.id, type=PDF_TYPE, author=author, release_date=release_date)
        add_and_refresh(pdf)

    @staticmethod
    def _read_pdf_meta(pdf_file_path, name):
        pdf_stream = open(pdf_file_path, "rb")
        pdf_reader = PdfFileReader(pdf_stream)
        pdf_info = pdf_reader.getDocumentInfo()
        title = name
        author = None
        release_date = None
        if '/Title' in pdf_info:
            title = pdf_info['/Title']
        if '/Author' in pdf_info:
            author = pdf_info['/Author']
        if '/CreationDate' in pdf_info:
            date_string = pdf_info['/CreationDate'].replace("'", '')        # D:YYYYMMDDHHmmSSOHH'mm'
            release_date = datetime.strptime(date_string, 'D:%Y%m%d%H%M%S')
        return title, author, release_date

<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Simple utility to calculate different types of hashes from given files/data
"""

from __future__ import print_function, unicode_literals

import argparse
import fileinput
import hashlib
import os
import sys
import zlib

try:
    from io import StringIO
except ImportError:
    from StringIO import StringIO

__version__ = '1.0.0'

MAX_INPUT_READ = 4*1024**2
DEFAULT_ALGORITHM = 'sha1'
E_OK = 0
E_FAIL = 1
E_FAIL_NO_FILES = 2
VERIFICATION_OK = 'OK'
VERIFICATION_FAIL = 'FAILED'

# pylint: disable=missing-docstring
def _get_available_hash_algorithms():
    try:
        available = hashlib.algorithms_available
    except AttributeError:
        try:
            available = set(hashlib.algorithms)
        except AttributeError:
            available = {'md5', 'sha1', 'sha224', 'sha256', 'sha384', 'sha512'}

    aliases = {
        'sha1': {'DSA', 'DSA-SHA', 'dsaEncryption', 'dsaWithSHA', 'ecdsa-with-SHA1'}
    }

    for name in aliases:
        if name in available:
            available -= aliases[name]

    to_remove = set()
    for algo in available:
        lowered = algo.lower()
        if algo != lowered and lowered in available:
            to_remove.add(algo)

    available -= to_remove
    return available

AVAILABLE_HASH_ALGORITHMS = _get_available_hash_algorithms()
AVAILABLE_CHECKSUM_ALGORITHMS = {'crc32', 'adler32'}
AVAILABLE_ALGORITHMS = sorted(AVAILABLE_HASH_ALGORITHMS | AVAILABLE_CHECKSUM_ALGORITHMS)

def checksum_file(file_path, algo, max_input_read=4*1024**2):
    """
    Calculate checksum
    :param file_path: path or '-' for STDIN
    :param algo:
    :param max_input_read:
    :return:
    """
    if algo not in AVAILABLE_CHECKSUM_ALGORITHMS:
        raise ValueError("Unknown algorithm: %s" % algo)

    value = 0
    hasher = getattr(zlib, algo)

    fh = open(file_path, 'rb') if file_path != '-' else sys.stdin
    while True:
        data = fh.read(max_input_read)
        if not len(data):
            break

        value = hasher(data, value)

    if file_path != '-':
        fh.close()

    value = value & 0xffffffff
    return hex(value)[2:]

def hash_file(file_path, algo, max_input_read=4*1024**2):
    """
    Calculate hash
    :param file_path: path or '-' for STDIN
    :param algo:
    :param max_input_read:
    :return:
    """
    hasher = hashlib.new(algo)

    fh = open(file_path, 'r') if file_path != '-' else sys.stdin
    while True:
        data = fh.read(max_input_read)
        if not len(data):
            break

        data = data.encode()
        hasher.update(data)

    if file_path != '-':
        fh.close()

    return hasher.hexdigest()

def _get_file_helpers():
    helpers = {algo: hash_file for algo in AVAILABLE_HASH_ALGORITHMS}
    helpers.update({algo: checksum_file for algo in AVAILABLE_CHECKSUM_ALGORITHMS})
    return helpers

FILE_HELPERS = _get_file_helpers()

OPENED_FILES = {'succes': 0, 'fail': 0}
def fileinput_openhook_safe(name, mode):
    try:
        fh = open(name, mode)
    except Exception as exc:
        OPENED_FILES['fail'] += 1
        print('%s: cannot open (%s)' % (name, exc.args[1]))
        return StringIO()
    else:
        OPENED_FILES['succes'] += 1
        return fh

def mode_generate_algo_symlinks(args):
    """
    Print commmands to create symlinks for hashfile for every known algorithm
    :return:
    """
    current = os.path.realpath(sys.argv[0])
    bindir = os.path.dirname(current)
    for algo in AVAILABLE_ALGORITHMS:
        print('ln -s %s %s' % (current, os.path.join(bindir, algo)))

    return E_OK

def mode_calculate(args):
    """
    Calculate hases and pront them to stdout
    :param args:
    :return:
    """
    for i, filename in enumerate(args.files):
        algo = args.algorithm[i] if len(args.algorithm) > i else args.algorithm[-1]

        file_helper = FILE_HELPERS[algo]
        try:
            filehash = file_helper(filename, algo=algo, max_input_read=args.max_input_read)
        except (OSError, IOError) as exc:
            print('ERROR: %s %s' % (filename, str(exc)), file=sys.stderr)
        else:
            print('%s: %s %s' % (algo, filehash, filename))

    return E_OK

def mode_check(args):
    """
    Verify calculated checksums
    :param args:
    :return:
    """
    if not args.files or args.files == ['-']:
        print('ERROR: no files to check specified', file=sys.stderr)
        sys.exit(1)

    def _quiet(inp1, inp2):
        if inp1 != inp2:
            print('%s: %s' % (filename, VERIFICATION_FAIL))
            return
        return True

    def _verbose(inp1, inp2):
        if inp1 == inp2:
            print('%s: %s' % (filename, VERIFICATION_OK))
        else:
            print('%s: %s' % (filename, VERIFICATION_FAIL))

        return True

    def _status(inp1, inp2):
        if inp1 != inp2:
            sys.exit(E_FAIL)

        return True

    if args.quiet:
        verifier = _quiet
    elif args.status:
        verifier = _status
    else:
        verifier = _verbose

    exit_code = E_OK

    for line in fileinput.input(args.files, openhook=fileinput_openhook_safe):
        line = line.strip()
        if not line or line.startswith('#'):
            continue

        try:
            algo, expected_filehash = line.split(': ', 1)
            expected_filehash, filename = expected_filehash.split(' ')
        except ValueError:
            if args.warn:
                print('ERROR: %s Incorrect format' % (fileinput.filename()), file=sys.stderr)
            continue

        file_helper = FILE_HELPERS[algo]
        try:
            filehash = file_helper(filename, algo=algo, max_input_read=args.max_input_read)
        except (OSError, IOError) as exc:
            print('ERROR: %s %s' % (filename, str(exc)), file=sys.stderr)
        else:
            if not verifier(filehash, expected_filehash):
                exit_code = E_FAIL
    if exit_code == E_OK and OPENED_FILES['fail'] > 0:
        exit_code = E_FAIL_NO_FILES
    return exit_code

def mode_default(args):
    """
    Do nothing
    :param args:
    :return:
    """
    return E_OK

def parse_args(argv):
    """
    Parse input params
    :param argv:
    :return:
    """
    parser = argparse.ArgumentParser(description='Calculate hash of some files',
        epilog='Algorithm can be also set from program name (for example call program as sha1 to use sha1 algorithm)')
    parser.add_argument('--algorithm', '-a', default=[], action='append', choices=AVAILABLE_ALGORITHMS,
        help='algorithm used to calculate hash '
             'If given more then one, then use different algorithms for different files (use first algo to first '
             'file, second algo to second file etc. If there is more files then algorithms, last algorithm from '
             'list is used.')
    parser.add_argument('--generate-algo-symlinks', action='store_const', dest='mode', const='generate-algo-symlinks',
        help='Show aliases for every algorithm handled by hashfile')
    # http://linux.die.net/man/1/md5sum
    parser.add_argument('--check', '-c', action='store_const', dest='mode', const='check',
        help='')
    parser.add_argument('--quiet', '-q', action='store_true',
        help='')
    parser.add_argument('--status', '-s', action='store_true',
        help='')
    parser.add_argument('--warn', '-w', action='store_true',
        help='')
    parser.add_argument('--max-input-read', default=MAX_INPUT_READ,
        help='maximum data size for read at once')
    parser.add_argument('files', metavar='file', type=str, nargs='*',
        help='list of files (stdin by default)')

    parser.set_defaults(mode='calculate')

    args = parser.parse_args()

    if args.mode != 'check' and (args.quiet or args.status or args.warn):
        parser.error('--quiet, --status and --warn options are available only with --check option')

    if len(args.algorithm) == 0:
        args.algorithm = [DEFAULT_ALGORITHM, ]
    elif os.path.basename(sys.argv[0]) in AVAILABLE_ALGORITHMS:
        args.algorithm = [os.path.basename(sys.argv[0]), ]

    if not args.files:
        args.files = ['-']
    # PYTHON2
    elif hasattr(args.files[0], 'decode'):
        args.files = [filename.decode('utf-8') for filename in args.files]

    return args

# pylint: disable=missing-docstring
def main():
    args = parse_args(sys.argv[1:])

    modes = {
        'calculate': mode_calculate,
        'check': mode_check,
        'generate-algo-symlinks': mode_generate_algo_symlinks,
    }
    handler = modes.get(args.mode, mode_default)
    exit_code = handler(args)

    sys.exit(exit_code)

if __name__ == '__main__':
    main()
<commit_msg>added some help<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Simple utility to calculate different types of hashes from given files/data
"""

from __future__ import print_function, unicode_literals

import argparse
import fileinput
import hashlib
import os
import sys
import zlib

try:
    from io import StringIO
except ImportError:
    from StringIO import StringIO

__version__ = '1.0.0'

MAX_INPUT_READ = 4*1024**2
DEFAULT_ALGORITHM = 'sha1'
E_OK = 0
E_FAIL = 1
E_FAIL_NO_FILES = 2
VERIFICATION_OK = 'OK'
VERIFICATION_FAIL = 'FAILED'

# pylint: disable=missing-docstring
def _get_available_hash_algorithms():
    try:
        available = hashlib.algorithms_available
    except AttributeError:
        try:
            available = set(hashlib.algorithms)
        except AttributeError:
            available = {'md5', 'sha1', 'sha224', 'sha256', 'sha384', 'sha512'}

    aliases = {
        'sha1': {'DSA', 'DSA-SHA', 'dsaEncryption', 'dsaWithSHA', 'ecdsa-with-SHA1'}
    }

    for name in aliases:
        if name in available:
            available -= aliases[name]

    to_remove = set()
    for algo in available:
        lowered = algo.lower()
        if algo != lowered and lowered in available:
            to_remove.add(algo)

    available -= to_remove
    return available

AVAILABLE_HASH_ALGORITHMS = _get_available_hash_algorithms()
AVAILABLE_CHECKSUM_ALGORITHMS = {'crc32', 'adler32'}
AVAILABLE_ALGORITHMS = sorted(AVAILABLE_HASH_ALGORITHMS | AVAILABLE_CHECKSUM_ALGORITHMS)

def checksum_file(file_path, algo, max_input_read=4*1024**2):
    """
    Calculate checksum
    :param file_path: path or '-' for STDIN
    :param algo:
    :param max_input_read:
    :return:
    """
    if algo not in AVAILABLE_CHECKSUM_ALGORITHMS:
        raise ValueError("Unknown algorithm: %s" % algo)

    value = 0
    hasher = getattr(zlib, algo)

    fh = open(file_path, 'rb') if file_path != '-' else sys.stdin
    while True:
        data = fh.read(max_input_read)
        if not len(data):
            break

        value = hasher(data, value)

    if file_path != '-':
        fh.close()

    value = value & 0xffffffff
    return hex(value)[2:]

def hash_file(file_path, algo, max_input_read=4*1024**2):
    """
    Calculate hash
    :param file_path: path or '-' for STDIN
    :param algo:
    :param max_input_read:
    :return:
    """
    hasher = hashlib.new(algo)

    fh = open(file_path, 'r') if file_path != '-' else sys.stdin
    while True:
        data = fh.read(max_input_read)
        if not len(data):
            break

        data = data.encode()
        hasher.update(data)

    if file_path != '-':
        fh.close()

    return hasher.hexdigest()

def _get_file_helpers():
    helpers = {algo: hash_file for algo in AVAILABLE_HASH_ALGORITHMS}
    helpers.update({algo: checksum_file for algo in AVAILABLE_CHECKSUM_ALGORITHMS})
    return helpers

FILE_HELPERS = _get_file_helpers()

OPENED_FILES = {'succes': 0, 'fail': 0}
def fileinput_openhook_safe(name, mode):
    try:
        fh = open(name, mode)
    except Exception as exc:
        OPENED_FILES['fail'] += 1
        print('%s: cannot open (%s)' % (name, exc.args[1]))
        return StringIO()
    else:
        OPENED_FILES['succes'] += 1
        return fh

def mode_generate_algo_symlinks(args):
    """
    Print commmands to create symlinks for hashfile for every known algorithm
    :return:
    """
    current = os.path.realpath(sys.argv[0])
    bindir = os.path.dirname(current)
    for algo in AVAILABLE_ALGORITHMS:
        print('ln -s %s %s' % (current, os.path.join(bindir, algo)))

    return E_OK

def mode_calculate(args):
    """
    Calculate hases and pront them to stdout
    :param args:
    :return:
    """
    for i, filename in enumerate(args.files):
        algo = args.algorithm[i] if len(args.algorithm) > i else args.algorithm[-1]

        file_helper = FILE_HELPERS[algo]
        try:
            filehash = file_helper(filename, algo=algo, max_input_read=args.max_input_read)
        except (OSError, IOError) as exc:
            print('ERROR: %s %s' % (filename, str(exc)), file=sys.stderr)
        else:
            print('%s: %s %s' % (algo, filehash, filename))

    return E_OK

def mode_check(args):
    """
    Verify calculated checksums
    :param args:
    :return:
    """
    if not args.files or args.files == ['-']:
        print('ERROR: no files to check specified', file=sys.stderr)
        sys.exit(1)

    def _quiet(inp1, inp2):
        if inp1 != inp2:
            print('%s: %s' % (filename, VERIFICATION_FAIL))
            return
        return True

    def _verbose(inp1, inp2):
        if inp1 == inp2:
            print('%s: %s' % (filename, VERIFICATION_OK))
        else:
            print('%s: %s' % (filename, VERIFICATION_FAIL))

        return True

    def _status(inp1, inp2):
        if inp1 != inp2:
            sys.exit(E_FAIL)

        return True

    if args.quiet:
        verifier = _quiet
    elif args.status:
        verifier = _status
    else:
        verifier = _verbose

    exit_code = E_OK

    for line in fileinput.input(args.files, openhook=fileinput_openhook_safe):
        line = line.strip()
        if not line or line.startswith('#'):
            continue

        try:
            algo, expected_filehash = line.split(': ', 1)
            expected_filehash, filename = expected_filehash.split(' ')
        except ValueError:
            if args.warn:
                print('ERROR: %s Incorrect format' % (fileinput.filename()), file=sys.stderr)
            continue

        file_helper = FILE_HELPERS[algo]
        try:
            filehash = file_helper(filename, algo=algo, max_input_read=args.max_input_read)
        except (OSError, IOError) as exc:
            print('ERROR: %s %s' % (filename, str(exc)), file=sys.stderr)
        else:
            if not verifier(filehash, expected_filehash):
                exit_code = E_FAIL
    if exit_code == E_OK and OPENED_FILES['fail'] > 0:
        exit_code = E_FAIL_NO_FILES
    return exit_code

def mode_default(args):
    """
    Do nothing
    :param args:
    :return:
    """
    return E_OK

def parse_args(argv):
    """
    Parse input params
    :param argv:
    :return:
    """
    parser = argparse.ArgumentParser(description='Calculate hash of some files',
        epilog='Algorithm can be also set from program name (for example call program as sha1 to use sha1 algorithm)')
    parser.add_argument('--algorithm', '-a', default=[], action='append', choices=AVAILABLE_ALGORITHMS,
        help='algorithm used to calculate hash '
             'If given more then one, then use different algorithms for different files (use first algo to first '
             'file, second algo to second file etc. If there is more files then algorithms, last algorithm from '
             'list is used.')
    parser.add_argument('--generate-algo-symlinks', action='store_const', dest='mode', const='generate-algo-symlinks',
        help='Show aliases for every algorithm handled by hashfile')
    # http://linux.die.net/man/1/md5sum
    parser.add_argument('--check', '-c', action='store_const', dest='mode', const='check',
        help='read checksums from the FILEs and check them ')
    parser.add_argument('--quiet', '-q', action='store_true',
        help='don\'t print OK for each successfully verified file')
    parser.add_argument('--status', '-s', action='store_true',
        help='don\'t output anything, status code shows success')
    parser.add_argument('--warn', '-w', action='store_true',
        help='warn about improperly formatted checksum lines')
    parser.add_argument('--max-input-read', default=MAX_INPUT_READ,
        help='maximum data size for read at once')
    parser.add_argument('files', metavar='FILE', type=str, nargs='*',
        help='list of files (stdin by default)')

    parser.set_defaults(mode='calculate')

    args = parser.parse_args()

    if args.mode != 'check' and (args.quiet or args.status or args.warn):
        parser.error('--quiet, --status and --warn options are available only with --check option')

    if len(args.algorithm) == 0:
        args.algorithm = [DEFAULT_ALGORITHM, ]
    elif os.path.basename(sys.argv[0]) in AVAILABLE_ALGORITHMS:
        args.algorithm = [os.path.basename(sys.argv[0]), ]

    if not args.files:
        args.files = ['-']
    # PYTHON2
    elif hasattr(args.files[0], 'decode'):
        args.files = [filename.decode('utf-8') for filename in args.files]

    return args

# pylint: disable=missing-docstring
def main():
    args = parse_args(sys.argv[1:])

    modes = {
        'calculate': mode_calculate,
        'check': mode_check,
        'generate-algo-symlinks': mode_generate_algo_symlinks,
    }
    handler = modes.get(args.mode, mode_default)
    exit_code = handler(args)

    sys.exit(exit_code)

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


<commit_before>"""
Main API for Vega-lite spec generation
"""

try:
    import traitlets as T
except ImportError:
    from IPython.utils import traitlets as T

from .utils import parse_shorthand

class BaseObject(T.HasTraits):

    def __contains__(self, key):
        value = getattr(self, key)
        return (value is not None) and (not (not isinstance(value, bool) and not value))

    def to_dict(self):
        result = {}
        for k in self.traits():
            if k in self:
                v = getattr(self, k)
                if isinstance(BaseObject):
                    result[k] = v.to_dict()
                else:
                    result[k] = v

class Data(BaseObject):

    formatType = T.Enum(['json','csv'], default_value='json')
    url = T.Unicode('', allow_none=True)
    data = T.List([], allow_none=True)

class Scale(BaseObject):
    pass

class Axis(BaseObject):
    pass

class Band(BaseObject):
    pass

class Legend(BaseObject):
    pass

class SortItems(BaseObject):
    name = T.Unicode(allow_none=True)
    aggregate = T.Enum(['avg','sum','min','max','count'], allow_none=True)
    reverse = T.Bool(False)

class Position(BaseObject):

    name = T.Unicode('')
    type = T.Enum(['N','O','Q','T'], allow_none=True)
    aggregate = T.Enum(['avg','sum','median','min','max','count'], default_value=None)
    timeUnit = T.Enum(['year','month','day','date','hours','minutes','seconds'], allow_none=True)
    bin = T.Union([T.Bool(),T.Int()], default_value=False)
    scale = T.Instance(Scale, allow_none=True)
    axis = T.Instance(Axis, allow_none=True)
    band = T.Instance(Band, allow_none=True)
    sort = T.List(T.Instance(SortItems), allow_none=True)

class Index(BaseObject):

    name = T.Unicode('')
    type = T.Enum(['N','O','Q','T'], allow_none=True)
    timeUnit = T.Enum(['year','month','day','date','hours','minutes','seconds'], allow_none=True)
    bin = T.Union([T.Bool(),T.Int()], default_value=False)
    aggregate = T.Enum(['count'], allow_none=True)
    padding = T.CFloat(0.1)
    sort = T.List(T.Instance(SortItems), allow_none=True)
    axis = T.Instance(Axis, allow_none=True)
    height = T.CInt(150)

class Size(BaseObject):
    name = T.Unicode('')
    type = T.Enum(['N','O','Q','T'])
    aggregate = T.Enum(['avg','sum','median','min','max','count'])
    timeUnit = T.Enum(['year','month','day','date','hours','minutes','seconds'], allow_none=True)
    bin = T.Union([T.Bool(),T.Int()], default_value=False)
    scale = T.Instance(Scale, allow_none=True)
    legend = T.Instance(Legend, allow_none=True)
    value = T.CInt(30)
    sort = T.List(T.Instance(SortItems), allow_none=True)

class Color(BaseObject):
    name = T.Unicode('')
    type = T.Enum(['N','O','Q','T'])
    aggregate = T.Enum(['avg','sum','median','min','max','count'])
    timeUnit = T.Enum(['year','month','day','date','hours','minutes','seconds'], allow_none=True)
    bin = T.Union([T.Bool(),T.Int()], default_value=False)
    scale = T.Instance(Scale, allow_none=True)
    legend = T.Instance(Legend, allow_none=True)
    value = T.Unicode('#4682b4')
    opacity = T.Float(1.0)
    sort = T.List(T.Instance(SortItems), allow_none=True)

class Shape(BaseObject):
    name = T.Unicode('')
    type = T.Enum(['N','O','Q','T'])
    aggregate = T.Enum(['count'], allow_none=True)
    timeUnit = T.Enum(['year','month','day','date','hours','minutes','seconds'], allow_none=True)
    bin = T.Union([T.Bool(),T.Int()], default_value=False)
    legend = T.Instance(Legend, allow_none=True)
    value = T.Enum(['circle','square','cross','diamond','triangle-up','triangle-down'], default_value='circle')
    filled = T.Bool(False)
    sort = T.List(T.Instance(SortItems), allow_none=True)

class Encoding(BaseObject):

    x = T.Union([T.Instance(Position),T.Unicode()], allow_none=True)
    y = T.Union([T.Instance(Position),T.Unicode()], allow_none=True)
    row = T.Union([T.Instance(Index),T.Unicode()], allow_none=True)
    col = T.Union([T.Instance(Index),T.Unicode()], allow_none=True)
    size = T.Instance(Size, allow_none=True)
    color = T.Instance(Color, allow_none=True)
    shape = T.Instance(Shape, allow_none=True)
    parent = T.Instance('altair.api.Viz')

    def _x_changed(self, name, old, new):
        if isinstance(new, str):
            result = parse_shorthand(new)
            self.x = Position(**result)

    def _y_changed(self, name, old, new):
        if isinstance(new, str):
            result = parse_shorthand(new)
            self.y = Position(**result)

    def _row_changed(self, name, old, new):
        if isinstance(new, str):
            result = parse_shorthand(new)
            self.row = Index(**result)

    def _col_changed(self, name, old, new):
        if isinstance(new, str):
            result = parse_shorthand(new)
            self.col = Index(**result)

class Viz(BaseObject):

    marktype = T.Enum(['point','tick','bar','line',
                     'area','circle','square','text'], default_value='point')
    _data = T.Instance(Data, allow_none=True)
    data = T.Any()
    encoding = T.Instance(Encoding)
    # _encoding = T.Union([T.Instance(Encoding),T.Dict])

    # def __encoding_changed(self, name, old, new):
    #     if isinstance(new, dict):
    #         self._encoding = Encoding(**new)

    def __init__(self, data, **kwargs):
        kwargs['data'] = data
        super(Viz,self).__init__(self, **kwargs)

    def _data_changed(self, name, old, new):
        self.encoding.infer_types(new)

    def encode(self, **kwargs):
        self.encoding = Encoding(**kwargs)
        return self

    def mark(self, mt):
        self.marktype = mt
        return self

    def mark_point(self):
        return self.mark('point')

    def mark_tick(self):
        return self.mark('tick')

    def mark_bar(self):
        return self.mark('bar')

    def mark_line(self):
        return self.mark('line')

    def mark_area(self):
        return self.mark('area')

    def mark_circle(self):
        return self.mark('circle')

    def mark_square(self):
        return self.mark('square')

    def mark_text(self):
        return self.mark('text')
<commit_msg>Fiing None conventions<commit_after>"""
Main API for Vega-lite spec generation
"""

try:
    import traitlets as T
except ImportError:
    from IPython.utils import traitlets as T

from .utils import parse_shorthand

class BaseObject(T.HasTraits):

    def __contains__(self, key):
        value = getattr(self, key)
        return (value is not None) and (not (not isinstance(value, bool) and not value))

    def to_dict(self):
        result = {}
        for k in self.traits():
            if k in self:
                v = getattr(self, k)
                if isinstance(BaseObject):
                    result[k] = v.to_dict()
                else:
                    result[k] = v

class Data(BaseObject):

    formatType = T.Enum(['json','csv'], default_value='json')
    url = T.Unicode(default_value=None, allow_none=True)
    data = T.List(default_value=None, allow_none=True)

class Scale(BaseObject):
    pass

class Axis(BaseObject):
    pass

class Band(BaseObject):
    pass

class Legend(BaseObject):
    pass

class SortItems(BaseObject):
    name = T.Unicode(default_value=None, allow_none=True)
    aggregate = T.Enum(['avg','sum','min','max','count'], default_value=True)
    reverse = T.Bool(False)

class Position(BaseObject):

    name = T.Unicode('')
    type = T.Enum(['N','O','Q','T'], default_value=None, allow_none=True)
    aggregate = T.Enum(['avg','sum','median','min','max','count'], default_value=None, allow_none=True)
    timeUnit = T.Enum(['year','month','day','date','hours','minutes','seconds'], default_value=None, allow_none=True)
    bin = T.Union([T.Bool(),T.Int()], default_value=False)
    scale = T.Instance(Scale, default_value=None, allow_none=True)
    axis = T.Instance(Axis, default_value=None, allow_none=True)
    band = T.Instance(Band, default_value=None, allow_none=True)
    sort = T.List(T.Instance(SortItems), defalut_value=None, allow_none=True)

class Index(BaseObject):

    name = T.Unicode(default_value=None, allow_none=True)
    type = T.Enum(['N','O','Q','T'], default_value=None, allow_none=True)
    timeUnit = T.Enum(['year','month','day','date','hours','minutes','seconds'], default_value=None, allow_none=True)
    bin = T.Union([T.Bool(),T.Int()], default_value=False)
    aggregate = T.Enum(['count'], default_value=None, allow_none=True)
    padding = T.CFloat(0.1)
    sort = T.List(T.Instance(SortItems), default_value=None, allow_none=True)
    axis = T.Instance(Axis, default_value=None, allow_none=True)
    height = T.CInt(150)

class Size(BaseObject):
    name = T.Unicode(default_value=None, allow_none=True)
    type = T.Enum(['N','O','Q','T'], default_value=None, allow_none=True)
    aggregate = T.Enum(['avg','sum','median','min','max','count'], default_value=None, allow_none=True)
    timeUnit = T.Enum(['year','month','day','date','hours','minutes','seconds'], default_value=None, allow_none=True)
    bin = T.Union([T.Bool(),T.Int()], default_value=False)
    scale = T.Instance(Scale, default_value=None, allow_none=True)
    legend = T.Instance(Legend, default_value=None, allow_none=True)
    value = T.CInt(30)
    sort = T.List(T.Instance(SortItems), default_value=None, allow_none=True)

class Color(BaseObject):
    name = T.Unicode(default_value=None, allow_none=True)
    type = T.Enum(['N','O','Q','T'], default_value=None, allow_none=True)
    aggregate = T.Enum(['avg','sum','median','min','max','count'], default_value=None, allow_none=True)
    timeUnit = T.Enum(['year','month','day','date','hours','minutes','seconds'], default_value=None, allow_none=True)
    bin = T.Union([T.Bool(),T.Int()], default_value=False)
    scale = T.Instance(Scale, default_value=None, allow_none=True)
    legend = T.Instance(Legend, default_value=None, allow_none=True)
    value = T.Unicode('#4682b4')
    opacity = T.Float(1.0)
    sort = T.List(T.Instance(SortItems), default_value=None, allow_none=True)

class Shape(BaseObject):
    name = T.Unicode(default_value=None, allow_none=True)
    type = T.Enum(['N','O','Q','T'], default_value=None, allow_none=True)
    aggregate = T.Enum(['count'], default_value=None, allow_none=True)
    timeUnit = T.Enum(['year','month','day','date','hours','minutes','seconds'], default_value=None, allow_none=True)
    bin = T.Union([T.Bool(),T.Int()], default_value=False)
    legend = T.Instance(Legend, default_value=None, allow_none=True)
    value = T.Enum(['circle','square','cross','diamond','triangle-up','triangle-down'], default_value='circle')
    filled = T.Bool(False)
    sort = T.List(T.Instance(SortItems), default_value=None, allow_none=True)

class Encoding(BaseObject):

    x = T.Union([T.Instance(Position),T.Unicode()], default_value=None, allow_none=True)
    y = T.Union([T.Instance(Position),T.Unicode()], default_value=None, allow_none=True)
    row = T.Union([T.Instance(Index),T.Unicode()], default_value=None, allow_none=True)
    col = T.Union([T.Instance(Index),T.Unicode()], default_value=None, allow_none=True)
    size = T.Instance(Size, default_value=None, allow_none=True)
    color = T.Instance(Color, default_value=None, allow_none=True)
    shape = T.Instance(Shape, default_value=None, allow_none=True)

    def _x_changed(self, name, old, new):
        if isinstance(new, str):
            result = parse_shorthand(new)
            self.x = Position(**result)

    def _y_changed(self, name, old, new):
        if isinstance(new, str):
            result = parse_shorthand(new)
            self.y = Position(**result)

    def _row_changed(self, name, old, new):
        if isinstance(new, str):
            result = parse_shorthand(new)
            self.row = Index(**result)

    def _col_changed(self, name, old, new):
        if isinstance(new, str):
            result = parse_shorthand(new)
            self.col = Index(**result)

class Viz(BaseObject):

    marktype = T.Enum(['point','tick','bar','line',
                     'area','circle','square','text'], default_value='point')
    _data = T.Instance(Data, default_value=None, allow_none=True)
    data = T.Any(default_value=None, allow_none=True)
    encoding = T.Instance(Encoding, default_value=None, allow_none=True)
    # _encoding = T.Union([T.Instance(Encoding),T.Dict])

    # def __encoding_changed(self, name, old, new):
    #     if isinstance(new, dict):
    #         self._encoding = Encoding(**new)

    def __init__(self, data, **kwargs):
        kwargs['data'] = data
        super(Viz,self).__init__(self, **kwargs)

    def encode(self, **kwargs):
        self.encoding = Encoding(**kwargs)
        return self

    def mark(self, mt):
        self.marktype = mt
        return self

    def mark_point(self):
        return self.mark('point')

    def mark_tick(self):
        return self.mark('tick')

    def mark_bar(self):
        return self.mark('bar')

    def mark_line(self):
        return self.mark('line')

    def mark_area(self):
        return self.mark('area')

    def mark_circle(self):
        return self.mark('circle')

    def mark_square(self):
        return self.mark('square')

    def mark_text(self):
        return self.mark('text')
<|endoftext|>


<commit_before>import psycopg2
import operator
from math import sqrt,fabs,floor

def distance_between (x1,y1,x2,y2):
    return fabs(sqrt( (x2 - x1)**2 + (y2 - y1)**2 ))

def getlocation(cur,version,userstopcode):
    cur = conn.cursor()
    cur.execute("""SELECT locationx_ew::integer,locationy_ns::integer FROM point WHERE version = %s and pointcode = %s""",[version,userstopcode])
    return cur.fetchone()

def set_beginend(conn,pool,userstopcodebegin,userstopcodeend,distance):
    version = pool[0][0]
    daow = pool[0][1]
    linkvalidfrom = pool[0][9]
    if pool[0][4] != userstopcodebegin:
        found = False
        for row in pool:
            if row[4] == userstopcodebegin:
                pool.remove(row)
                pool.insert(0,row)
                found = True
                break
        if found is False:
            x,y = getlocation(conn,version,userstopcodebegin)
            first = [version,daow,userstopcodebegin,userstopcodeend,userstopcodebegin,0,x,y,0]
            pool.insert(0,first)
            cur = conn.cursor()
            cur.execute("""INSERT INTO POOL VALUES ('POOL',%s,'I','VTN',%s,%s,%s,'VTN',%s,0,NULL,NULL,NULL)""",
                           [version,userstopcodebegin,userstopcodeend,linkvalidfrom,userstopcodebegin])
    if pool[-1][4] != userstopcodeend:
        found = False
        for row in pool:
            if row[4] == userstopcodeend:
                pool.remove(row)
                pool.append(row)
                found = True
                break
        if found is False:
            x,y = getlocation(conn,version,userstopcodeend)
            last = [version,daow,userstopcodebegin,userstopcodeend,userstopcodeend,distance,x,y,distance]
            pool.append(last)
            cur = conn.cursor()
            cur.execute("""INSERT INTO POOL VALUES ('POOL',%s,'I','VTN',%s,%s,%s,'VTN',%s,%s,NULL,NULL,NULL)""",
                           [version,userstopcodebegin,userstopcodeend,linkvalidfrom,userstopcodeend,distance])
            print pool

def fix_pool(conn):
    cur = conn.cursor()
    cur.execute("""SELECT version,userstopcodebegin,userstopcodeend,distance::integer from link""")
    for version,userstopcodebegin,userstopcodeend,distance in cur.fetchall():
        print (version,userstopcodebegin,userstopcodeend,distance)
        cur.execute("""
SELECT
pl.version, --0
pl.dataownercode, --1
userstopcodebegin, --2
userstopcodeend, --3
pl.pointcode, --4
distancesincestartoflink::integer, --5
pt.locationx_ew::integer, --6
pt.locationy_ns::integer, --7
ST_DISTANCE(st_makepoint(pt.locationx_ew,pt.locationy_ns),st_makepoint(ut.locationx_ew,ut.locationy_ns))::integer, --8
linkvalidfrom --9
FROM pool as pl, point as pt,point as ut
WHERE
pl.pointdataownercode = pt.dataownercode AND
pl.version = pt.version AND
pl.pointcode = pt.pointcode AND
pl.version = ut.version AND
pl.userstopcodebegin = ut.pointcode AND
pl.userstopcodebegin = %s AND
pl.userstopcodeend = %s AND
pl.version = %s""",[userstopcodebegin,userstopcodeend,version])
        pool = [list(x) for x in cur.fetchall()]
        pool[0][1] = 3
        if len(pool) == 2:
            if pool[0][4] == userstopcodebegin and pool[1][4] == userstopcodebegin:
                if pool[1][5] != distance:
                    raise Exception('2-pool but not right distance')
                continue
        pool = sorted(pool, key=operator.itemgetter(8))
        set_beginend(conn,pool,userstopcodebegin,userstopcodeend,distance)
        clean_pool = [pool[0]]
        del(pool[0])
        while len(pool) > 1:
            pointer = 0
            min_dist = 999999991
            for i in range(len(pool)):
                dist = distance_between(pool[i][6],pool[i][7],clean_pool[-1][6],clean_pool[-1][7])
                if pointer is None or (dist < min_dist and pool[i][3] != pool[i][4]):
                    pointer = i
                    min_dist = dist
            pool[pointer][5] = int(floor(min_dist + clean_pool[-1][5]))
            clean_pool.append(pool[pointer])
            pool.remove(pool[pointer])
        clean_pool.append(pool[-1])
        if clean_pool[-1][5] < clean_pool[-2][5]:
            for j in range(len(clean_pool)):
                i = len(clean_pool)-j-1
                if i > 0 and clean_pool[i][5] < clean_pool[i-1][5]:
                    clean_pool[i-1][5] = clean_pool[i][5]-1
            print clean_pool
        for p in clean_pool[1:-1]:
            cur.execute("update pool set distancesincestartoflink = %s where version = %s and userstopcodebegin = %s and userstopcodeend = %s and pointcode = %s",[p[5],p[0],p[2],p[3],p[4]])
<commit_msg>Veolia shape fixer: don't commit<commit_after>import psycopg2
import operator
from math import sqrt,fabs,floor

def distance_between (x1,y1,x2,y2):
    return fabs(sqrt( (x2 - x1)**2 + (y2 - y1)**2 ))

def getlocation(conn,version,userstopcode):
    cur = conn.cursor()
    cur.execute("""SELECT locationx_ew::integer,locationy_ns::integer FROM point WHERE version = %s and pointcode = %s""",[version,userstopcode])
    return cur.fetchone()

def set_beginend(conn,pool,userstopcodebegin,userstopcodeend,distance):
    version = pool[0][0]
    daow = pool[0][1]
    linkvalidfrom = pool[0][9]
    if pool[0][4] != userstopcodebegin:
        found = False
        for row in pool:
            if row[4] == userstopcodebegin:
                pool.remove(row)
                pool.insert(0,row)
                found = True
                break
        if found is False:
            x,y = getlocation(conn,version,userstopcodebegin)
            first = [version,daow,userstopcodebegin,userstopcodeend,userstopcodebegin,0,x,y,0]
            pool.insert(0,first)
            cur = conn.cursor()
            cur.execute("""INSERT INTO POOL VALUES ('POOL',%s,'I','VTN',%s,%s,%s,'VTN',%s,0,NULL,NULL,NULL)""",
                           [version,userstopcodebegin,userstopcodeend,linkvalidfrom,userstopcodebegin])
    if pool[-1][4] != userstopcodeend:
        found = False
        for row in pool:
            if row[4] == userstopcodeend:
                pool.remove(row)
                pool.append(row)
                found = True
                break
        if found is False:
            x,y = getlocation(conn,version,userstopcodeend)
            last = [version,daow,userstopcodebegin,userstopcodeend,userstopcodeend,distance,x,y,distance]
            pool.append(last)
            cur = conn.cursor()
            cur.execute("""INSERT INTO POOL VALUES ('POOL',%s,'I','VTN',%s,%s,%s,'VTN',%s,%s,NULL,NULL,NULL)""",
                           [version,userstopcodebegin,userstopcodeend,linkvalidfrom,userstopcodeend,distance])
            print pool

def fix_pool(conn):
    cur = conn.cursor()
    cur.execute("""SELECT version,userstopcodebegin,userstopcodeend,distance::integer from link""")
    for version,userstopcodebegin,userstopcodeend,distance in cur.fetchall():
        print (version,userstopcodebegin,userstopcodeend,distance)
        cur.execute("""
SELECT
pl.version, --0
pl.dataownercode, --1
userstopcodebegin, --2
userstopcodeend, --3
pl.pointcode, --4
distancesincestartoflink::integer, --5
pt.locationx_ew::integer, --6
pt.locationy_ns::integer, --7
ST_DISTANCE(st_makepoint(pt.locationx_ew,pt.locationy_ns),st_makepoint(ut.locationx_ew,ut.locationy_ns))::integer, --8
linkvalidfrom --9
FROM pool as pl, point as pt,point as ut
WHERE
pl.pointdataownercode = pt.dataownercode AND
pl.version = pt.version AND
pl.pointcode = pt.pointcode AND
pl.version = ut.version AND
pl.userstopcodebegin = ut.pointcode AND
pl.userstopcodebegin = %s AND
pl.userstopcodeend = %s AND
pl.version = %s""",[userstopcodebegin,userstopcodeend,version])
        pool = [list(x) for x in cur.fetchall()]
        pool[0][1] = 3
        if len(pool) == 2:
            if pool[0][4] == userstopcodebegin and pool[1][4] == userstopcodebegin:
                if pool[1][5] != distance:
                    raise Exception('2-pool but not right distance')
                continue
        pool = sorted(pool, key=operator.itemgetter(8))
        set_beginend(conn,pool,userstopcodebegin,userstopcodeend,distance)
        clean_pool = [pool[0]]
        del(pool[0])
        while len(pool) > 1:
            pointer = 0
            min_dist = 999999991
            for i in range(len(pool)):
                dist = distance_between(pool[i][6],pool[i][7],clean_pool[-1][6],clean_pool[-1][7])
                if pointer is None or (dist < min_dist and pool[i][3] != pool[i][4]):
                    pointer = i
                    min_dist = dist
            pool[pointer][5] = int(floor(min_dist + clean_pool[-1][5]))
            clean_pool.append(pool[pointer])
            pool.remove(pool[pointer])
        clean_pool.append(pool[-1])
        if clean_pool[-1][5] < clean_pool[-2][5]:
            for j in range(len(clean_pool)):
                i = len(clean_pool)-j-1
                if i > 0 and clean_pool[i][5] < clean_pool[i-1][5]:
                    clean_pool[i-1][5] = clean_pool[i][5]-1
            print clean_pool
        for p in clean_pool[1:-1]:
            cur.execute("update pool set distancesincestartoflink = %s where version = %s and userstopcodebegin = %s and userstopcodeend = %s and pointcode = %s",[p[5],p[0],p[2],p[3],p[4]])
<|endoftext|>