<s> '''<STR_LIT>''' <EOL> import os <EOL> import re <EOL> import threading <EOL> import mmap <EOL> import sys <EOL> from collections import defaultdict <EOL> import struct <EOL> import marshal <EOL> try : <EOL> import cPickle as pickle <EOL> except ImportError : <EOL> import pickle <EOL> from cola . core . utils import iterable <EOL> from cola . core . mq . utils import labelize <EOL> class StoreExistsError ( Exception ) : pass <EOL> class StoreNotSafetyShutdown ( Exception ) : pass <EOL> class StoreNoSpaceForPut ( Exception ) : pass <EOL> STORE_FILE_SIZE = <NUM_LIT:4> * <NUM_LIT> * <NUM_LIT> <EOL> LEGAL_STORE_FILE_REGEX = re . compile ( "<STR_LIT>" ) <EOL> READ_ENTRANCE , WRITE_ENTRANCE = range ( <NUM_LIT:2> ) <EOL> MARSHAL , PICKLE = '<STR_LIT:m>' , '<STR_LIT:p>' <EOL> class Store ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , working_dir , size = STORE_FILE_SIZE , <EOL> deduper = None , mkdirs = False , <EOL> create_lock_file = False ) : <EOL> """<STR_LIT>""" <EOL> self . lock = threading . Lock ( ) <EOL> self . store_file_size = size <EOL> self . deduper = deduper <EOL> self . dir_ = working_dir <EOL> if mkdirs and not os . path . exists ( self . dir_ ) : <EOL> os . makedirs ( self . dir_ ) <EOL> self . create_lock_file = create_lock_file <EOL> self . lock_file = os . path . join ( self . dir_ , '<STR_LIT>' ) <EOL> with self . lock : <EOL> if create_lock_file is True : <EOL> if os . path . exists ( self . lock_file ) : <EOL> raise StoreExistsError ( '<STR_LIT>' ) <EOL> else : <EOL> open ( self . lock_file , '<STR_LIT:w>' ) . close ( ) <EOL> self . stopped = False <EOL> self . inited = False <EOL> self . legal_files = [ ] <EOL> self . file_handles = defaultdict ( lambda : None ) <EOL> self . map_handles = defaultdict ( lambda : None ) <EOL> def shutdown ( self ) : <EOL> if self . stopped : return <EOL> self . stopped = True <EOL> try : <EOL> for handle in self . map_handles . values ( ) : <EOL> if handle is not None : <EOL> handle . close ( ) <EOL> for handle in self . file_handles . values ( ) : <EOL> if handle is not None : <EOL> handle . close ( ) <EOL> finally : <EOL> with self . lock : <EOL> if self . create_lock_file is True : <EOL> os . remove ( self . lock_file ) <EOL> self . inited = False <EOL> def init ( self ) : <EOL> with self . lock : <EOL> if self . inited : return <EOL> files = os . listdir ( self . dir_ ) <EOL> for fi in files : <EOL> if fi == '<STR_LIT>' : continue <EOL> file_path = os . path . join ( self . dir_ , fi ) <EOL> if not os . path . isfile ( file_path ) or LEGAL_STORE_FILE_REGEX . match ( fi ) is None : <EOL> raise StoreNotSafetyShutdown ( '<STR_LIT>' ) <EOL> else : <EOL> self . legal_files . append ( file_path ) <EOL> self . legal_files = sorted ( self . legal_files , key = lambda k : int ( os . path . basename ( k ) ) ) <EOL> if len ( self . legal_files ) > <NUM_LIT:0> : <EOL> read_file_handle = self . file_handles [ READ_ENTRANCE ] = open ( self . legal_files [ - <NUM_LIT:1> ] , '<STR_LIT>' ) <EOL> self . map_handles [ READ_ENTRANCE ] = mmap . mmap ( read_file_handle . fileno ( ) , <EOL> self . store_file_size ) <EOL> if len ( self . legal_files ) == <NUM_LIT:1> : <EOL> self . file_handles [ WRITE_ENTRANCE ] = self . file_handles [ READ_ENTRANCE ] <EOL> self . map_handles [ WRITE_ENTRANCE ] = self . map_handles [ READ_ENTRANCE ] <EOL> else : <EOL> write_file_handle = self . file_handles [ WRITE_ENTRANCE ] = open ( self . legal_files [ <NUM_LIT:0> ] , '<STR_LIT>' ) <EOL> self . map_handles [ WRITE_ENTRANCE ] = mmap . mmap ( write_file_handle . fileno ( ) , <EOL> self . store_file_size ) <EOL> self . inited = True <EOL> def _generate_file ( self ) : <EOL> prev = None <EOL> if len ( self . legal_files ) > <NUM_LIT:0> : <EOL> fn = os . path . basename ( self . legal_files [ <NUM_LIT:0> ] ) <EOL> prev = int ( LEGAL_STORE_FILE_REGEX . match ( fn ) . group ( ) ) <EOL> current = str ( prev - <NUM_LIT:1> if prev is not None else sys . maxint ) <EOL> file_path = os . path . join ( self . dir_ , current ) <EOL> if len ( self . legal_files ) > <NUM_LIT:1> : <EOL> self . map_handles [ WRITE_ENTRANCE ] . close ( ) <EOL> self . file_handles [ WRITE_ENTRANCE ] . close ( ) <EOL> self . legal_files . insert ( <NUM_LIT:0> , file_path ) <EOL> open ( file_path , '<STR_LIT:w>' ) . close ( ) <EOL> write_file_handle = self . file_handles [ WRITE_ENTRANCE ] = open ( file_path , '<STR_LIT>' ) <EOL> write_file_handle . write ( '<STR_LIT:\x00>' * self . store_file_size ) <EOL> write_file_handle . flush ( ) <EOL> self . map_handles [ WRITE_ENTRANCE ] = mmap . mmap ( write_file_handle . fileno ( ) , <EOL> self . store_file_size ) <EOL> if len ( self . legal_files ) == <NUM_LIT:1> : <EOL> self . map_handles [ READ_ENTRANCE ] = self . map_handles [ WRITE_ENTRANCE ] <EOL> self . file_handles [ READ_ENTRANCE ] = self . file_handles [ WRITE_ENTRANCE ] <EOL> def _destroy_file ( self ) : <EOL> if len ( self . legal_files ) == <NUM_LIT:0> : <EOL> return <EOL> self . map_handles [ READ_ENTRANCE ] . close ( ) <EOL> self . file_handles [ READ_ENTRANCE ] . close ( ) <EOL> if len ( self . legal_files ) == <NUM_LIT:1> : <EOL> self . map_handles [ WRITE_ENTRANCE ] . close ( ) <EOL> self . file_handles [ WRITE_ENTRANCE ] . close ( ) <EOL> self . map_handles . clear ( ) <EOL> self . file_handles . clear ( ) <EOL> elif len ( self . legal_files ) == <NUM_LIT:2> : <EOL> self . map_handles [ READ_ENTRANCE ] = self . map_handles [ WRITE_ENTRANCE ] <EOL> self . file_handles [ READ_ENTRANCE ] = self . file_handles [ WRITE_ENTRANCE ] <EOL> else : <EOL> read_file_handle = self . file_handles [ READ_ENTRANCE ] = open ( self . legal_files [ - <NUM_LIT:2> ] , '<STR_LIT>' ) <EOL> self . map_handles [ READ_ENTRANCE ] = mmap . mmap ( read_file_handle . fileno ( ) , <EOL> self . store_file_size ) <EOL> self . legal_files . pop ( - <NUM_LIT:1> ) <EOL> def _stringfy ( self , obj ) : <EOL> try : <EOL> return MARSHAL + marshal . dumps ( obj ) <EOL> except ValueError : <EOL> return PICKLE + pickle . dumps ( obj ) <EOL> def _destringfy ( self , src_str ) : <EOL> if len ( src_str ) < <NUM_LIT:2> : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> t , str_ = src_str [ <NUM_LIT:0> ] , src_str [ <NUM_LIT:1> : ] <EOL> if t == MARSHAL : <EOL> obj = marshal . loads ( str_ ) <EOL> elif t == PICKLE : <EOL> obj = pickle . loads ( str_ ) <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> return obj <EOL> def _seek_writable_pos ( self , map_handle ) : <EOL> pos = <NUM_LIT:0> <EOL> while True : <EOL> if pos + <NUM_LIT:4> <= self . store_file_size : <EOL> size , = struct . unpack ( '<STR_LIT:I>' , map_handle [ pos : pos + <NUM_LIT:4> ] ) <EOL> if size == <NUM_LIT:0> : <EOL> return pos <EOL> pos += ( <NUM_LIT:4> + size ) <EOL> else : <EOL> return - <NUM_LIT:1> <EOL> return - <NUM_LIT:1> <EOL> def put_one ( self , obj , force = False , commit = True ) : <EOL> """<STR_LIT>""" <EOL> if self . stopped : return <EOL> if isinstance ( obj , str ) and obj . strip ( ) == '<STR_LIT>' : <EOL> return <EOL> if not force and self . deduper is not None : <EOL> prop = labelize ( obj ) <EOL> if self . deduper . exist ( prop ) : <EOL> return <EOL> if len ( self . legal_files ) == <NUM_LIT:0> : <EOL> self . _generate_file ( ) <EOL> obj_str = self . _stringfy ( obj ) <EOL> if len ( obj_str ) + <NUM_LIT:4> > self . store_file_size : <EOL> raise StoreNoSpaceForPut ( '<STR_LIT>' ) <EOL> with self . lock : <EOL> m = self . map_handles [ WRITE_ENTRANCE ] <EOL> pos = self . _seek_writable_pos ( m ) <EOL> size = pos + <NUM_LIT:4> + len ( obj_str ) <EOL> if pos < <NUM_LIT:0> or size > self . store_file_size : <EOL> m . flush ( ) <EOL> self . _generate_file ( ) <EOL> pos = <NUM_LIT:0> <EOL> size = <NUM_LIT:4> + len ( obj_str ) <EOL> m = self . map_handles [ WRITE_ENTRANCE ] <EOL> m [ : size ] = m [ : pos ] + struct . pack ( '<STR_LIT:I>' , len ( obj_str ) ) + obj_str <EOL> if commit is True : <EOL> m . flush ( ) <EOL> return obj <EOL> def put ( self , objects , force = False , commit = True ) : <EOL> if self . stopped : return <EOL> self . init ( ) <EOL> if isinstance ( objects , basestring ) or not iterable ( objects ) : <EOL> return self . put_one ( objects , force , commit ) <EOL> remains = [ ] <EOL> for obj in objects : <EOL> result = self . put_one ( obj , force , commit = False ) <EOL> if result is not None : <EOL> remains . append ( result ) <EOL> m = self . map_handles [ WRITE_ENTRANCE ] <EOL> if len ( remains ) > <NUM_LIT:0> and m is not None : <EOL> with self . lock : <EOL> m . flush ( ) <EOL> return remains <EOL> def get_one ( self , commit = True ) : <EOL> """<STR_LIT>""" <EOL> if self . stopped : return <EOL> self . init ( ) <EOL> m = self . map_handles [ READ_ENTRANCE ] <EOL> while m is not None : <EOL> with self . lock : <EOL> size , = struct . unpack ( '<STR_LIT:I>' , m [ : <NUM_LIT:4> ] ) <EOL> if size == <NUM_LIT:0> : <EOL> self . _destroy_file ( ) <EOL> m = self . map_handles [ READ_ENTRANCE ] <EOL> else : <EOL> obj = self . _destringfy ( m [ <NUM_LIT:4> : <NUM_LIT:4> + size ] ) <EOL> m [ : ] = m [ <NUM_LIT:4> + size : ] + '<STR_LIT:\x00>' * ( <NUM_LIT:4> + size ) <EOL> if commit is True : <EOL> m . flush ( ) <EOL> return obj <EOL> def get ( self , size = <NUM_LIT:1> ) : <EOL> if size <= <NUM_LIT:1> : <EOL> return self . get_one ( ) <EOL> self . init ( ) <EOL> results = [ ] <EOL> for _ in range ( size ) : <EOL> obj = self . get_one ( ) <EOL> if obj is not None : <EOL> results . append ( obj ) <EOL> m = self . map_handles [ READ_ENTRANCE ] <EOL> if m is not None : <EOL> with self . lock : <EOL> m . flush ( ) <EOL> return results <EOL> def __enter__ ( self ) : <EOL> return self <EOL> def __exit__ ( self , type_ , value , traceback ) : <EOL> self . shutdown ( ) </s>
