datastructures.py 1.85 KB
Newer Older
Antoine Millet's avatar
Antoine Millet committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21
#!/usr/bin/env python
#coding:utf8

import threading

class BytesBuffer(object):
    '''
    BytesBuffers objects are useful to create socket buffers.
    
    Its behavior is pretty simple and looks like a FIFO queue: you can push data
    on it with push method, and pull a specified amount of data.
    
    .. note::
       All calls to :class:`BytesBuffer` objects are thread-safe.
    '''
    
    def __init__(self, data=None):
        if data is None:
            self._buf = ''
        else:
            self._buf = data
22
        self._lock = threading.RLock()
Antoine Millet's avatar
Antoine Millet committed
23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65
    
    def push(self, data):
        '''
        Push data on buffer.
        
        :param data: the data to push on the buffer
        :type data: :class:`str` object
        '''
        
        self._lock.acquire()
        self._buf += data
        self._lock.release()
    
    def pull(self, size=0):
        '''
        Pull the specified amount of data on the buffer.
        
        If size is equal to 0, will return the entire content. If buffer 
        contains less of data than asked, :meth:`pull` will return all the 
        available data.
        
        :param size: size (in bytes) of data to pull on the buffer
        :type size: 0 or positive integer
        :return: asked data
        :rtype: :class:`str` object
        '''
        
        self._lock.acquire()
        if size == 0:
            buf = self._buf
            self._buf = ''
        else:
            buf = self._buf[:size]
            self._buf = self._buf[size:]
        self._lock.release()
        return buf
    
    def __len__(self):
        '''
        Return the size of the buffer.
        '''
        
        return len(self._buf)
66 67 68 69

    def __enter__(self):
        return self._lock.__enter__()

70 71
    def __exit__(self, exc_type, exc_value, traceback):
        return self._lock.__exit__(exc_type, exc_value, traceback)