<s> import pickle <EOL> from huey . exceptions import QueueException <EOL> class TaskRegistry ( object ) : <EOL> """<STR_LIT>""" <EOL> _ignore = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> _registry = { } <EOL> _periodic_tasks = [ ] <EOL> def task_to_string ( self , task ) : <EOL> return '<STR_LIT:%s>' % ( task . __name__ ) <EOL> def register ( self , task_class ) : <EOL> klass_str = self . task_to_string ( task_class ) <EOL> if klass_str in self . _ignore : <EOL> return <EOL> if klass_str not in self . _registry : <EOL> self . _registry [ klass_str ] = task_class <EOL> if hasattr ( task_class , '<STR_LIT>' ) : <EOL> self . _periodic_tasks . append ( task_class ( ) ) <EOL> def unregister ( self , task_class ) : <EOL> klass_str = self . task_to_string ( task_class ) <EOL> if klass_str in self . _registry : <EOL> del ( self . _registry [ klass_str ] ) <EOL> for task in self . _periodic_tasks : <EOL> if isinstance ( task , task_class ) : <EOL> self . _periodic_tasks . remove ( task ) <EOL> def __contains__ ( self , klass_str ) : <EOL> return klass_str in self . _registry <EOL> def get_message_for_task ( self , task ) : <EOL> """<STR_LIT>""" <EOL> return pickle . dumps ( ( <EOL> task . task_id , <EOL> self . task_to_string ( type ( task ) ) , <EOL> task . execute_time , <EOL> task . retries , <EOL> task . retry_delay , <EOL> task . get_data ( ) , <EOL> ) ) <EOL> def get_task_class ( self , klass_str ) : <EOL> klass = self . _registry . get ( klass_str ) <EOL> if not klass : <EOL> raise QueueException ( '<STR_LIT>' % klass_str ) <EOL> return klass <EOL> def get_task_for_message ( self , msg ) : <EOL> """<STR_LIT>""" <EOL> raw = pickle . loads ( msg ) <EOL> task_id , klass_str , execute_time , retries , delay , data = raw <EOL> klass = self . get_task_class ( klass_str ) <EOL> return klass ( data , task_id , execute_time , retries , delay ) <EOL> def get_periodic_tasks ( self ) : <EOL> return self . _periodic_tasks <EOL> registry = TaskRegistry ( ) </s>
