utils.py

# -*- coding: utf-8 -*-

import os, sys, subprocess
from socket import gethostname, gethostbyname
from threading import Lock
from logging import debug


def enum(*sequential, **named):
    '''
    '''
    enums = dict(zip(sequential, range(len(sequential))), **named)
    return type('Enum', (), enums)


class Enum(dict):
    '''
    Usage example:
    
    $ states = Enum(SLEEPING, CONNECTED="Conn", DISCONNECTED="Disco")
    $ print states.CONNECTED
    1
    $ states.SLEEPING
    0
    $ states.CONNECTED_str
    "Conn"
    $ states.SLEEPING_str
    None
    $ 0 in states
    True
    $ 10 in states
    False
    $ s = "Conn"
    $ if s in states:
    $     val = states[s]
    '''
    def __init__(self, *named, **valued):
        '''
        '''
        # FIXME not memory optimal structures
        self._items_map = {}
        self._items_rmap = {}
        self._items_val = []
        toks = [(name, None) for name in named]
        toks.extend([i for i in valued.iteritems()])
        for key, tok in enumerate(toks):
            self._items_map[tok[0]] = key
            self._items_rmap[key] = tok[0]
            self._items_val.append(tok[1])
    
    def __contains__(self, item):
        '''
        '''
        if isinstance(item, int):
            return item in self._items_rmap
        else:
            return item in self._items_val
    
    def __getitem__(self, key):
        '''
        '''
        if isinstance(key, int):
            return self._items_val[key]
        else:
            for idx, val in enumerate(self._items_val):
                if val == key:
                    return idx
            raise IndexError('object not found in enum')
    
    def __getattr__(self, name):
        '''
        '''
        if name in self._items_map:
            return self._items_map[name]
        else:
            if name[-4:] == '_str':
                n = name[:-4]
                if n in self._items_map:
                    return self._items_val[self._items_map[n]]
        raise AttributeError('not an enum member')

    def get(self, key, default=None):
        '''
        '''
        try:
            return self[key]
        except IndexError as err:
            return default


class RWLock(object):
    '''
    '''
    def __init__(self):
        '''
        '''
        self._mutex = Lock()
        self._writemutex = Lock()
        self._readers = 0
        self._writers = 0
        self.read = self._RLock(self)
        self.write = self._WLock(self)
    
    
    class _Lock(object):
        '''
        '''
        def __init__(self, rwlock):
            '''
            '''
            self._parent = rwlock
    
    
    class _WLock(_Lock):
        '''
        '''
        def __enter__(self):
            '''
            '''
            self.acquire()
        
        def __exit__(self, exc_type, exc_value, traceback):
            '''
            '''
            self.release()
        
        def acquire(self):
            '''
            '''
            with self._parent._mutex:
                self._parent._writers += 1
            self._parent._writemutex.acquire()
        
        def release(self):
            '''
            '''
            with self._parent._mutex:
                self._parent._writers -= 1
            self._parent._writemutex.release()
    
    
    class _RLock(_Lock):
        '''
        '''
        def __enter__(self):
            '''
            '''
            self.acquire()
        
        def __exit__(self, exc_type, exc_value, traceback):
            '''
            '''
            self.release()
        
        def acquire(self):
            '''
            '''
            self._parent._mutex.acquire()
            if self._parent._writers > 0 or self._parent._readers == 0:
                self._parent._mutex.release()
                self._parent._writemutex.acquire()
                self._parent._mutex.acquire()
            self._parent._readers += 1
            self._parent._mutex.release()
        
        def release(self):
            '''
            '''
            self._parent._mutex.acquire()
            self._parent._readers -= 1
            if self._parent._readers == 0:
                self._parent._writemutex.release()
            self._parent._mutex.release()


class Exec(object):
    '''
    '''
    @staticmethod
    def call(cmd, shell=False, input=None, capture_output=True,
                                                            merge_output=True):
        '''
        '''
        if capture_output:
            stdout = subprocess.PIPE
            if merge_output:
                stderr = subprocess.STDOUT
            else:
                stderr = subprocess.PIPE
        else:
            stdout = open(os.devnull, 'w')
            stderr = stdout
        debug('Exec.call: cmd=`%s`', ' '.join(cmd))
        if input is not None:
            stdin = subprocess.PIPE
            debug('Exec.call: input=`%s`', input)
        else:
            stdin = None
        proc = subprocess.Popen(cmd, shell=shell, stdout=stdout, stderr=stderr,
                                                                stdin=stdin)
        output = proc.communicate(input=input)
        debug('Exec.call: rc=`%d` out=`%s`', proc.returncode, output)
        return (proc.returncode, output)
    
    @staticmethod
    def silent(cmd, **kwargs):
        '''
        '''
        kwargs['capture_output'] = False
        return Exec.call(cmd, **kwargs)[0]