pvl/irker.py
author Tero Marttila <terom@paivola.fi>
Thu, 10 Jan 2013 20:03:44 +0200
changeset 79 530c2aa73a97
parent 72 7bb07131c2b5
permissions -rw-r--r--
pvl.irker: refactor
import optparse, sys

import logging; log = logging.getLogger('pvl.irker')

# proto
import socket, json

def parser (parser, connect='tcp://localhost/', target=None) :
    """
        Optparse option group.
    """

    irker = optparse.OptionGroup(parser, 'Irker output')
    
    irker.add_option('--irker', metavar='URL',  default=connect,
            help="Irker daemon URL")

    return irker

def apply (options) :
    """
        Return Irker (XXX: target) from options.
    """
    
    # None -> stdout
    return Irker(options.irker)

def connect (host=None, port=None, family=socket.AF_UNSPEC, socktype=socket.SOCK_STREAM) :
    """
        Return a TCP/UDP socket connected to the given host/port using getaddrinfo.

        TODO: timeout?
    """

    log.debug("%s:%s: %s/%s", host, port, family, socktype)
    
    if host :
        flags = socket.AI_CANONNAME
    else :
        flags = 0

    addrinfo = socket.getaddrinfo(host, port, family, socktype, 0, flags)

    if not addrinfo :
        raise Exception("getaddrinfo: %s:%s: no results" % (host, port))

    for af, st, proto, name, addr in addrinfo :
        try :
            s = socket.socket(af, st, proto)

        except socket.error as error :
            log.warning("%s:%s: socket: %s", host, port, error)
            continue
        
        try :
            s.connect(addr)

        except socket.error as error :
            log.warning("%s:%s: connect: %s", host, port, error)
            continue
        
        log.info("%s", name)
        
        return s

    else :
        raise Exception("Unable to connect: %s:%s: %s" % (host, port, error))

import urlparse

class Irk (object) :
    """
        Irker JSON connection.

        TODO: timeout?
    """

    PORT = 6659

    SCHEME = {
        'tcp':  (socket.AF_INET, socket.SOCK_STREAM),
        'udp':  (socket.AF_INET, socket.SOCK_DGRAM),
        'unix': (socket.AF_UNIX, socket.SOCK_DGRAM),
    }

    @classmethod
    def connect (cls, url) :
        """
            Connect to given urllib URL, or None -> stdout
        """

        if not url :
            return cls(sys.stdout)

        family, socktype = cls.SCHEME[url.scheme]
        
        if family == socket.AF_UNIX :
            raise Exception("unix:// is not supported")
        else :
            # inet
            sock = connect(url.hostname, url.port or cls.PORT, family=family, socktype=socktype)
        
        return cls(sock.makefile('w'))

    def __init__ (self, file) :
        """
            Use given file-like object for output.
        """

        self.file = file

        log.debug("%s", file)

    def send (self, **opts) :
        """
            Raises IOError on write errors.
        """

        log.debug("%s", opts)
        
        # write line + flush
        json.dump(opts, self.file)
        self.file.write('\n')
        self.file.flush()

class IrkerTarget (object) :
    """
        A channel on an Irk connection.
    """

    def __init__ (self, irker, target) :
        self.irker = irker
        self.target = target

    def join (self) :
        log.info("%s", self)
        self.irker.send(to=str(self), privmsg='')

    def privmsg (self, *args) :
        for arg in args :
            log.info("%s: %s", self, arg)
            self.irker.send(to=str(self), privmsg=arg)

    __call__ = privmsg

    def __str__ (self) :
        return self.target

class Irker (object) :
    """
        Reconnecting irker.
    """

    def __init__ (self, url=None) :
        if url :
            self.url = urlparse.urlparse(url)
        else :
            self.url = None

        self.targets = {}
        
        self.connect()
    
    def connect (self) :
        """
            Connect, and fix up our targets.
        """

        self.irk = Irk.connect(self.url)

        # rejoin
        for target in self.targets.itervalues() :
            target.join()
    
    def send (self, **opts) :
        """
            Send on current irker connection.

            TODO: handle errors and reconnect?
        """

        self.irk.send(**opts)

    def target (self, target) :
        """
            Bind to given target URL, returning an IrkerTarget for sending messages.
        """

        if target not in self.targets :
            self.targets[target] = IrkerTarget(self, target)
            self.targets[target].join()
            
        return self.targets[target]

    __getitem__ = target