<s> """<STR_LIT>""" <EOL> import hmac <EOL> import logging <EOL> import os <EOL> import pprint <EOL> import uuid <EOL> from datetime import datetime <EOL> import collections <EOL> try : <EOL> import pickle <EOL> pickle = cPickle <EOL> except : <EOL> cPickle = None <EOL> import pickle <EOL> import zmq <EOL> from zmq . utils import jsonapi <EOL> from zmq . eventloop . ioloop import IOLoop <EOL> from zmq . eventloop . zmqstream import ZMQStream <EOL> from IPython . config . configurable import Configurable , LoggingConfigurable <EOL> from IPython . utils . importstring import import_item <EOL> from IPython . utils . jsonutil import extract_dates , squash_dates , date_default <EOL> from IPython . utils . traitlets import ( CBytes , Unicode , Bool , Any , Instance , Set , <EOL> DottedObjectName ) <EOL> def squash_unicode ( obj ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( obj , dict ) : <EOL> for key in list ( obj . keys ( ) ) : <EOL> obj [ key ] = squash_unicode ( obj [ key ] ) <EOL> if isinstance ( key , str ) : <EOL> obj [ squash_unicode ( key ) ] = obj . pop ( key ) <EOL> elif isinstance ( obj , list ) : <EOL> for i , v in enumerate ( obj ) : <EOL> obj [ i ] = squash_unicode ( v ) <EOL> elif isinstance ( obj , str ) : <EOL> obj = obj . encode ( '<STR_LIT:utf8>' ) <EOL> return obj <EOL> key = '<STR_LIT>' if jsonapi . jsonmod . __name__ == '<STR_LIT>' else '<STR_LIT:default>' <EOL> json_packer = lambda obj : jsonapi . dumps ( obj , ** { key : date_default } ) <EOL> json_unpacker = lambda s : extract_dates ( jsonapi . loads ( s ) ) <EOL> pickle_packer = lambda o : pickle . dumps ( o , - <NUM_LIT:1> ) <EOL> pickle_unpacker = pickle . loads <EOL> default_packer = json_packer <EOL> default_unpacker = json_unpacker <EOL> DELIM = b"<STR_LIT>" <EOL> class SessionFactory ( LoggingConfigurable ) : <EOL> """<STR_LIT>""" <EOL> logname = Unicode ( '<STR_LIT>' ) <EOL> def _logname_changed ( self , name , old , new ) : <EOL> self . log = logging . getLogger ( new ) <EOL> context = Instance ( '<STR_LIT>' ) <EOL> def _context_default ( self ) : <EOL> return zmq . Context . instance ( ) <EOL> session = Instance ( '<STR_LIT>' ) <EOL> loop = Instance ( '<STR_LIT>' , allow_none = False ) <EOL> def _loop_default ( self ) : <EOL> return IOLoop . instance ( ) <EOL> def __init__ ( self , ** kwargs ) : <EOL> super ( SessionFactory , self ) . __init__ ( ** kwargs ) <EOL> if self . session is None : <EOL> self . session = Session ( ** kwargs ) <EOL> class Message ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , msg_dict ) : <EOL> dct = self . __dict__ <EOL> for k , v in dict ( msg_dict ) . items ( ) : <EOL> if isinstance ( v , dict ) : <EOL> v = Message ( v ) <EOL> dct [ k ] = v <EOL> def __iter__ ( self ) : <EOL> return iter ( self . __dict__ . items ( ) ) <EOL> def __repr__ ( self ) : <EOL> return repr ( self . __dict__ ) <EOL> def __str__ ( self ) : <EOL> return pprint . pformat ( self . __dict__ ) <EOL> def __contains__ ( self , k ) : <EOL> return k in self . __dict__ <EOL> def __getitem__ ( self , k ) : <EOL> return self . __dict__ [ k ] <EOL> def msg_header ( msg_id , msg_type , username , session ) : <EOL> date = datetime . now ( ) <EOL> return locals ( ) <EOL> def extract_header ( msg_or_header ) : <EOL> """<STR_LIT>""" <EOL> if not msg_or_header : <EOL> return { } <EOL> try : <EOL> h = msg_or_header [ '<STR_LIT>' ] <EOL> except KeyError : <EOL> try : <EOL> h = msg_or_header [ '<STR_LIT>' ] <EOL> except KeyError : <EOL> raise <EOL> else : <EOL> h = msg_or_header <EOL> if not isinstance ( h , dict ) : <EOL> h = dict ( h ) <EOL> return h <EOL> class Session ( Configurable ) : <EOL> """<STR_LIT>""" <EOL> debug = Bool ( False , config = True , help = """<STR_LIT>""" ) <EOL> packer = DottedObjectName ( '<STR_LIT>' , config = True , <EOL> help = """<STR_LIT>""" ) <EOL> def _packer_changed ( self , name , old , new ) : <EOL> if new . lower ( ) == '<STR_LIT>' : <EOL> self . pack = json_packer <EOL> self . unpack = json_unpacker <EOL> elif new . lower ( ) == '<STR_LIT>' : <EOL> self . pack = pickle_packer <EOL> self . unpack = pickle_unpacker <EOL> else : <EOL> self . pack = import_item ( str ( new ) ) <EOL> unpacker = DottedObjectName ( '<STR_LIT>' , config = True , <EOL> help = """<STR_LIT>""" ) <EOL> def _unpacker_changed ( self , name , old , new ) : <EOL> if new . lower ( ) == '<STR_LIT>' : <EOL> self . pack = json_packer <EOL> self . unpack = json_unpacker <EOL> elif new . lower ( ) == '<STR_LIT>' : <EOL> self . pack = pickle_packer <EOL> self . unpack = pickle_unpacker <EOL> else : <EOL> self . unpack = import_item ( str ( new ) ) <EOL> session = Unicode ( '<STR_LIT>' , config = True , <EOL> help = """<STR_LIT>""" ) <EOL> def _session_default ( self ) : <EOL> return str ( uuid . uuid4 ( ) ) <EOL> username = Unicode ( os . environ . get ( '<STR_LIT>' , '<STR_LIT:username>' ) , config = True , <EOL> help = """<STR_LIT>""" ) <EOL> key = CBytes ( b'<STR_LIT>' , config = True , <EOL> help = """<STR_LIT>""" ) <EOL> def _key_changed ( self , name , old , new ) : <EOL> if new : <EOL> self . auth = hmac . HMAC ( new ) <EOL> else : <EOL> self . auth = None <EOL> auth = Instance ( hmac . HMAC ) <EOL> digest_history = Set ( ) <EOL> keyfile = Unicode ( '<STR_LIT>' , config = True , <EOL> help = """<STR_LIT>""" ) <EOL> def _keyfile_changed ( self , name , old , new ) : <EOL> with open ( new , '<STR_LIT:rb>' ) as f : <EOL> self . key = f . read ( ) . strip ( ) <EOL> pack = Any ( default_packer ) <EOL> def _pack_changed ( self , name , old , new ) : <EOL> if not isinstance ( new , collections . Callable ) : <EOL> raise TypeError ( "<STR_LIT>" % type ( new ) ) <EOL> unpack = Any ( default_unpacker ) <EOL> def _unpack_changed ( self , name , old , new ) : <EOL> if not isinstance ( new , collections . Callable ) : <EOL> raise TypeError ( "<STR_LIT>" % type ( new ) ) <EOL> def __init__ ( self , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> super ( Session , self ) . __init__ ( ** kwargs ) <EOL> self . _check_packers ( ) <EOL> self . none = self . pack ( { } ) <EOL> @ property <EOL> def msg_id ( self ) : <EOL> """<STR_LIT>""" <EOL> return str ( uuid . uuid4 ( ) ) <EOL> def _check_packers ( self ) : <EOL> """<STR_LIT>""" <EOL> pack = self . pack <EOL> unpack = self . unpack <EOL> msg = dict ( a = [ <NUM_LIT:1> , '<STR_LIT>' ] ) <EOL> try : <EOL> packed = pack ( msg ) <EOL> except Exception : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if not isinstance ( packed , bytes ) : <EOL> raise ValueError ( "<STR_LIT>" % type ( packed ) ) <EOL> try : <EOL> unpacked = unpack ( packed ) <EOL> except Exception : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> msg = dict ( t = datetime . now ( ) ) <EOL> try : <EOL> unpacked = unpack ( pack ( msg ) ) <EOL> except Exception : <EOL> self . pack = lambda o : pack ( squash_dates ( o ) ) <EOL> self . unpack = lambda s : extract_dates ( unpack ( s ) ) <EOL> def msg_header ( self , msg_type ) : <EOL> return msg_header ( self . msg_id , msg_type , self . username , self . session ) <EOL> def msg ( self , msg_type , content = None , parent = None , subheader = None ) : <EOL> """<STR_LIT>""" <EOL> msg = { } <EOL> msg [ '<STR_LIT>' ] = self . msg_header ( msg_type ) <EOL> msg [ '<STR_LIT>' ] = msg [ '<STR_LIT>' ] [ '<STR_LIT>' ] <EOL> msg [ '<STR_LIT>' ] = { } if parent is None else extract_header ( parent ) <EOL> msg [ '<STR_LIT>' ] = msg_type <EOL> msg [ '<STR_LIT:content>' ] = { } if content is None else content <EOL> sub = { } if subheader is None else subheader <EOL> msg [ '<STR_LIT>' ] . update ( sub ) <EOL> return msg <EOL> def sign ( self , msg_list ) : <EOL> """<STR_LIT>""" <EOL> if self . auth is None : <EOL> return b'<STR_LIT>' <EOL> h = self . auth . copy ( ) <EOL> for m in msg_list : <EOL> h . update ( m ) <EOL> return h . hexdigest ( ) . encode ( '<STR_LIT:ascii>' ) <EOL> def serialize ( self , msg , ident = None ) : <EOL> """<STR_LIT>""" <EOL> content = msg . get ( '<STR_LIT:content>' , { } ) <EOL> if content is None : <EOL> content = self . none <EOL> elif isinstance ( content , dict ) : <EOL> content = self . pack ( content ) <EOL> elif isinstance ( content , bytes ) : <EOL> pass <EOL> elif isinstance ( content , str ) : <EOL> content = content . encode ( '<STR_LIT:utf8>' ) <EOL> else : <EOL> raise TypeError ( "<STR_LIT>" % type ( content ) ) <EOL> real_message = [ self . pack ( msg [ '<STR_LIT>' ] ) , <EOL> self . pack ( msg [ '<STR_LIT>' ] ) , <EOL> content <EOL> ] <EOL> to_send = [ ] <EOL> if isinstance ( ident , list ) : <EOL> to_send . extend ( ident ) <EOL> elif ident is not None : <EOL> to_send . append ( ident ) <EOL> to_send . append ( DELIM ) <EOL> signature = self . sign ( real_message ) <EOL> to_send . append ( signature ) <EOL> to_send . extend ( real_message ) <EOL> return to_send <EOL> def send ( self , stream , msg_or_type , content = None , parent = None , ident = None , <EOL> buffers = None , subheader = None , track = False ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( stream , ( zmq . Socket , ZMQStream ) ) : <EOL> raise TypeError ( "<STR_LIT>" % type ( stream ) ) <EOL> elif track and isinstance ( stream , ZMQStream ) : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> if isinstance ( msg_or_type , ( Message , dict ) ) : <EOL> msg = msg_or_type <EOL> else : <EOL> msg = self . msg ( msg_or_type , content , parent , subheader ) <EOL> buffers = [ ] if buffers is None else buffers <EOL> to_send = self . serialize ( msg , ident ) <EOL> flag = <NUM_LIT:0> <EOL> if buffers : <EOL> flag = zmq . SNDMORE <EOL> _track = False <EOL> else : <EOL> _track = track <EOL> if track : <EOL> tracker = stream . send_multipart ( to_send , flag , copy = False , track = _track ) <EOL> else : <EOL> tracker = stream . send_multipart ( to_send , flag , copy = False ) <EOL> for b in buffers [ : - <NUM_LIT:1> ] : <EOL> stream . send ( b , flag , copy = False ) <EOL> if buffers : <EOL> if track : <EOL> tracker = stream . send ( buffers [ - <NUM_LIT:1> ] , copy = False , track = track ) <EOL> else : <EOL> tracker = stream . send ( buffers [ - <NUM_LIT:1> ] , copy = False ) <EOL> if self . debug : <EOL> pprint . pprint ( msg ) <EOL> pprint . pprint ( to_send ) <EOL> pprint . pprint ( buffers ) <EOL> msg [ '<STR_LIT>' ] = tracker <EOL> return msg <EOL> def send_raw ( self , stream , msg_list , flags = <NUM_LIT:0> , copy = True , ident = None ) : <EOL> """<STR_LIT>""" <EOL> to_send = [ ] <EOL> if isinstance ( ident , bytes ) : <EOL> ident = [ ident ] <EOL> if ident is not None : <EOL> to_send . extend ( ident ) <EOL> to_send . append ( DELIM ) <EOL> to_send . append ( self . sign ( msg_list ) ) <EOL> to_send . extend ( msg_list ) <EOL> stream . send_multipart ( msg_list , flags , copy = copy ) <EOL> def recv ( self , socket , mode = zmq . NOBLOCK , content = True , copy = True ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( socket , ZMQStream ) : <EOL> socket = socket . socket <EOL> try : <EOL> msg_list = socket . recv_multipart ( mode ) <EOL> except zmq . ZMQError as e : <EOL> if e . errno == zmq . EAGAIN : <EOL> return None , None <EOL> else : <EOL> raise <EOL> idents , msg_list = self . feed_identities ( msg_list , copy ) <EOL> try : <EOL> return idents , self . unpack_message ( msg_list , content = content , copy = copy ) <EOL> except Exception as e : <EOL> print ( ( idents , msg_list ) ) <EOL> raise e <EOL> def feed_identities ( self , msg_list , copy = True ) : <EOL> """<STR_LIT>""" <EOL> if copy : <EOL> idx = msg_list . index ( DELIM ) <EOL> return msg_list [ : idx ] , msg_list [ idx + <NUM_LIT:1> : ] <EOL> else : <EOL> failed = True <EOL> for idx , m in enumerate ( msg_list ) : <EOL> if m . bytes == DELIM : <EOL> failed = False <EOL> break <EOL> if failed : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> idents , msg_list = msg_list [ : idx ] , msg_list [ idx + <NUM_LIT:1> : ] <EOL> return [ m . bytes for m in idents ] , msg_list <EOL> def unpack_message ( self , msg_list , content = True , copy = True ) : <EOL> """<STR_LIT>""" <EOL> minlen = <NUM_LIT:4> <EOL> message = { } <EOL> if not copy : <EOL> for i in range ( minlen ) : <EOL> msg_list [ i ] = msg_list [ i ] . bytes <EOL> if self . auth is not None : <EOL> signature = msg_list [ <NUM_LIT:0> ] <EOL> if signature in self . digest_history : <EOL> raise ValueError ( "<STR_LIT>" % signature ) <EOL> self . digest_history . add ( signature ) <EOL> check = self . sign ( msg_list [ <NUM_LIT:1> : <NUM_LIT:4> ] ) <EOL> if not signature == check : <EOL> raise ValueError ( "<STR_LIT>" % signature ) <EOL> if not len ( msg_list ) >= minlen : <EOL> raise TypeError ( "<STR_LIT>" % minlen ) <EOL> message [ '<STR_LIT>' ] = self . unpack ( msg_list [ <NUM_LIT:1> ] ) <EOL> message [ '<STR_LIT>' ] = message [ '<STR_LIT>' ] [ '<STR_LIT>' ] <EOL> message [ '<STR_LIT>' ] = self . unpack ( msg_list [ <NUM_LIT:2> ] ) <EOL> if content : <EOL> message [ '<STR_LIT:content>' ] = self . unpack ( msg_list [ <NUM_LIT:3> ] ) <EOL> else : <EOL> message [ '<STR_LIT:content>' ] = msg_list [ <NUM_LIT:3> ] <EOL> message [ '<STR_LIT>' ] = msg_list [ <NUM_LIT:4> : ] <EOL> return message <EOL> def test_msg2obj ( ) : <EOL> am = dict ( x = <NUM_LIT:1> ) <EOL> ao = Message ( am ) <EOL> assert ao . x == am [ '<STR_LIT:x>' ] <EOL> am [ '<STR_LIT:y>' ] = dict ( z = <NUM_LIT:1> ) <EOL> ao = Message ( am ) <EOL> assert ao . y . z == am [ '<STR_LIT:y>' ] [ '<STR_LIT:z>' ] <EOL> k1 , k2 = '<STR_LIT:y>' , '<STR_LIT:z>' <EOL> assert ao [ k1 ] [ k2 ] == am [ k1 ] [ k2 ] <EOL> am2 = dict ( ao ) <EOL> assert am [ '<STR_LIT:x>' ] == am2 [ '<STR_LIT:x>' ] <EOL> assert am [ '<STR_LIT:y>' ] [ '<STR_LIT:z>' ] == am2 [ '<STR_LIT:y>' ] [ '<STR_LIT:z>' ] </s>
