<s> import collections <EOL> import schema <EOL> class DictLookup ( object ) : <EOL> def __init__ ( self ) : <EOL> self . container = { } <EOL> def append ( self , entity , entity_pk , node_position ) : <EOL> if not self . container . get ( entity ) : <EOL> self . container [ entity ] = { } <EOL> self . container [ entity ] [ entity_pk ] = node_position <EOL> def lookup ( self , entity , entity_pk ) : <EOL> if self . container . get ( entity ) : <EOL> if self . container [ entity ] . get ( entity_pk ) : <EOL> return self . container [ entity ] [ entity_pk ] <EOL> def start_entity ( self , entity ) : <EOL> pass <EOL> def stop_entity ( self , entity ) : <EOL> pass <EOL> class Node ( object ) : <EOL> def __init__ ( self , record , entity ) : <EOL> if not isinstance ( record , dict ) : <EOL> raise TypeError <EOL> self . record = record <EOL> if not isinstance ( entity , schema . Entity ) or not entity . get_primary_key_field ( ) : <EOL> raise TypeError <EOL> self . entity = entity <EOL> self . properties = { } <EOL> for field in entity . fields : <EOL> if field . column and isinstance ( field . column , schema . Column ) : <EOL> self . properties [ field . name ] = record . get ( field . column . name ) <EOL> else : <EOL> self . properties [ field . name ] = field . value <EOL> self . entity_pk = int ( record . get ( entity . get_primary_key_field ( ) . column . name ) ) <EOL> self . properties [ '<STR_LIT>' ] = self . entity . name <EOL> self . properties [ '<STR_LIT>' ] = None <EOL> def get_dict ( self , fields ) : <EOL> return dict ( ( e , self . properties . get ( e , '<STR_LIT>' ) ) for e in fields ) <EOL> def save_node_id ( self , node_id ) : <EOL> self . properties [ '<STR_LIT>' ] = node_id <EOL> def get_node_id ( self ) : <EOL> return self . properties [ '<STR_LIT>' ] <EOL> def get_entity_pk ( self ) : <EOL> return self . entity_pk <EOL> def __repr__ ( self ) : <EOL> return "<STR_LIT>" % ( self . entity . name , self . entity_pk , str ( self . properties ) ) <EOL> class Relation ( object ) : <EOL> def __init__ ( self , start , end , properties ) : <EOL> self . start = start <EOL> self . end = end <EOL> self . properties = properties <EOL> self . start_fk = None <EOL> self . start_target_entity = None <EOL> def get_dict ( self , fields ) : <EOL> if self . end : <EOL> values = dict ( ( <EOL> ( '<STR_LIT:start>' , self . start ) , <EOL> ( '<STR_LIT:end>' , self . end ) , <EOL> ) ) <EOL> values . update ( self . properties ) <EOL> return dict ( ( e , values . get ( e , '<STR_LIT>' ) ) for e in fields ) <EOL> else : <EOL> return None <EOL> def set_deferred_resolution ( self , fk , target_entity ) : <EOL> self . start_fk = fk <EOL> self . start_target_entity = target_entity <EOL> def __repr__ ( self ) : <EOL> return "<STR_LIT>" % ( <EOL> self . start , self . property , self . end or <NUM_LIT:0> , ) <EOL> class NodeList ( object ) : <EOL> def __init__ ( self ) : <EOL> self . node_list_size = <NUM_LIT:0> <EOL> self . reverse_lookup = DictLookup ( ) <EOL> self . entity_fields = set ( ) <EOL> self . last_lookup = None <EOL> self . last_lookup_result = None <EOL> def get_all_fields ( self ) : <EOL> return self . entity_fields <EOL> def update_entity_fields ( self , record ) : <EOL> self . entity_fields . update ( record . iterkeys ( ) ) <EOL> def add_node ( self , node ) : <EOL> """<STR_LIT>""" <EOL> self . node_list_size += <NUM_LIT:1> <EOL> node . save_node_id ( self . node_list_size ) <EOL> node_pk = node . get_entity_pk ( ) <EOL> if node_pk : <EOL> self . reverse_lookup . append ( node . entity . name , node_pk , node . get_node_id ( ) ) <EOL> else : <EOL> print node , "<STR_LIT>" <EOL> raise RuntimeError <EOL> self . last_lookup , self . last_lookup_result = None , None <EOL> return node . get_node_id ( ) <EOL> def lookup_node_pos ( self , entity_name , entity_pk ) : <EOL> if self . last_lookup == ( entity_name , entity_pk ) : <EOL> return self . last_lookup_result <EOL> else : <EOL> result = self . reverse_lookup . lookup ( entity_name , entity_pk ) <EOL> self . last_lookup = ( entity_name , entity_pk ) <EOL> self . last_lookup_result = result <EOL> return result <EOL> def iter_nodes ( self ) : <EOL> while True : <EOL> try : <EOL> yield self . node_list . popleft ( ) <EOL> except : <EOL> break <EOL> class RelationList ( object ) : <EOL> def __init__ ( self ) : <EOL> self . relation_list = collections . deque ( ) <EOL> def add_relation ( self , relation ) : <EOL> self . relation_list . append ( relation ) <EOL> def iter_rels ( self ) : <EOL> while True : <EOL> try : <EOL> yield self . relation_list . popleft ( ) <EOL> except : <EOL> break </s>
