<s> import sys <EOL> import inspect <EOL> import logging <EOL> import wrapt <EOL> from . compat import contextlib , collections <EOL> from . errors import UnhandledHTTPRequestError <EOL> from . matchers import requests_match , uri , method <EOL> from . patch import CassettePatcherBuilder <EOL> from . persist import load_cassette , save_cassette <EOL> from . serializers import yamlserializer <EOL> from . util import partition_dict <EOL> log = logging . getLogger ( __name__ ) <EOL> class CassetteContextDecorator ( object ) : <EOL> """<STR_LIT>""" <EOL> _non_cassette_arguments = ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> @ classmethod <EOL> def from_args ( cls , cassette_class , ** kwargs ) : <EOL> return cls ( cassette_class , lambda : dict ( kwargs ) ) <EOL> def __init__ ( self , cls , args_getter ) : <EOL> self . cls = cls <EOL> self . _args_getter = args_getter <EOL> self . __finish = None <EOL> def _patch_generator ( self , cassette ) : <EOL> with contextlib . ExitStack ( ) as exit_stack : <EOL> for patcher in CassettePatcherBuilder ( cassette ) . build ( ) : <EOL> exit_stack . enter_context ( patcher ) <EOL> log_format = '<STR_LIT>' <EOL> log . debug ( log_format . format ( <EOL> action = "<STR_LIT>" , path = cassette . _path <EOL> ) ) <EOL> yield cassette <EOL> log . debug ( log_format . format ( <EOL> action = "<STR_LIT>" , path = cassette . _path <EOL> ) ) <EOL> cassette . _save ( ) <EOL> def __enter__ ( self ) : <EOL> assert self . __finish is None , "<STR_LIT>" <EOL> other_kwargs , cassette_kwargs = partition_dict ( <EOL> lambda key , _ : key in self . _non_cassette_arguments , <EOL> self . _args_getter ( ) <EOL> ) <EOL> if other_kwargs . get ( '<STR_LIT>' ) : <EOL> transformer = other_kwargs [ '<STR_LIT>' ] <EOL> cassette_kwargs [ '<STR_LIT:path>' ] = transformer ( cassette_kwargs [ '<STR_LIT:path>' ] ) <EOL> self . __finish = self . _patch_generator ( self . cls . load ( ** cassette_kwargs ) ) <EOL> return next ( self . __finish ) <EOL> def __exit__ ( self , * args ) : <EOL> next ( self . __finish , None ) <EOL> self . __finish = None <EOL> @ wrapt . decorator <EOL> def __call__ ( self , function , instance , args , kwargs ) : <EOL> args_getter = self . _build_args_getter_for_decorator ( function ) <EOL> return type ( self ) ( self . cls , args_getter ) . _execute_function ( <EOL> function , args , kwargs <EOL> ) <EOL> def _execute_function ( self , function , args , kwargs ) : <EOL> if inspect . isgeneratorfunction ( function ) : <EOL> handler = self . _handle_coroutine <EOL> else : <EOL> handler = self . _handle_function <EOL> return handler ( function , args , kwargs ) <EOL> def _handle_coroutine ( self , function , args , kwargs ) : <EOL> """<STR_LIT>""" <EOL> with self as cassette : <EOL> coroutine = self . __handle_function ( cassette , function , args , kwargs ) <EOL> to_yield = next ( coroutine ) <EOL> while True : <EOL> try : <EOL> to_send = yield to_yield <EOL> except Exception : <EOL> to_yield = coroutine . throw ( * sys . exc_info ( ) ) <EOL> else : <EOL> to_yield = coroutine . send ( to_send ) <EOL> def __handle_function ( self , cassette , function , args , kwargs ) : <EOL> if cassette . inject : <EOL> return function ( cassette , * args , ** kwargs ) <EOL> else : <EOL> return function ( * args , ** kwargs ) <EOL> def _handle_function ( self , function , args , kwargs ) : <EOL> with self as cassette : <EOL> return self . __handle_function ( cassette , function , args , kwargs ) <EOL> @ staticmethod <EOL> def get_function_name ( function ) : <EOL> return function . __name__ <EOL> def _build_args_getter_for_decorator ( self , function ) : <EOL> def new_args_getter ( ) : <EOL> kwargs = self . _args_getter ( ) <EOL> if '<STR_LIT:path>' not in kwargs : <EOL> name_generator = ( kwargs . get ( '<STR_LIT>' ) or <EOL> self . get_function_name ) <EOL> path = name_generator ( function ) <EOL> kwargs [ '<STR_LIT:path>' ] = path <EOL> return kwargs <EOL> return new_args_getter <EOL> class Cassette ( object ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def load ( cls , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> new_cassette = cls ( ** kwargs ) <EOL> new_cassette . _load ( ) <EOL> return new_cassette <EOL> @ classmethod <EOL> def use_arg_getter ( cls , arg_getter ) : <EOL> return CassetteContextDecorator ( cls , arg_getter ) <EOL> @ classmethod <EOL> def use ( cls , ** kwargs ) : <EOL> return CassetteContextDecorator . from_args ( cls , ** kwargs ) <EOL> def __init__ ( self , path , serializer = yamlserializer , record_mode = '<STR_LIT>' , <EOL> match_on = ( uri , method ) , before_record_request = None , <EOL> before_record_response = None , custom_patches = ( ) , <EOL> inject = False ) : <EOL> self . _path = path <EOL> self . _serializer = serializer <EOL> self . _match_on = match_on <EOL> self . _before_record_request = before_record_request or ( lambda x : x ) <EOL> self . _before_record_response = before_record_response or ( lambda x : x ) <EOL> self . inject = inject <EOL> self . record_mode = record_mode <EOL> self . custom_patches = custom_patches <EOL> self . data = [ ] <EOL> self . play_counts = collections . Counter ( ) <EOL> self . dirty = False <EOL> self . rewound = False <EOL> @ property <EOL> def play_count ( self ) : <EOL> return sum ( self . play_counts . values ( ) ) <EOL> @ property <EOL> def all_played ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . play_count == len ( self ) <EOL> @ property <EOL> def requests ( self ) : <EOL> return [ request for ( request , response ) in self . data ] <EOL> @ property <EOL> def responses ( self ) : <EOL> return [ response for ( request , response ) in self . data ] <EOL> @ property <EOL> def write_protected ( self ) : <EOL> return self . rewound and self . record_mode == '<STR_LIT>' or self . record_mode == '<STR_LIT:none>' <EOL> def append ( self , request , response ) : <EOL> """<STR_LIT>""" <EOL> request = self . _before_record_request ( request ) <EOL> if not request : <EOL> return <EOL> response = self . _before_record_response ( response ) <EOL> if response is None : <EOL> return <EOL> self . data . append ( ( request , response ) ) <EOL> self . dirty = True <EOL> def filter_request ( self , request ) : <EOL> return self . _before_record_request ( request ) <EOL> def _responses ( self , request ) : <EOL> """<STR_LIT>""" <EOL> request = self . _before_record_request ( request ) <EOL> for index , ( stored_request , response ) in enumerate ( self . data ) : <EOL> if requests_match ( request , stored_request , self . _match_on ) : <EOL> yield index , response <EOL> def can_play_response_for ( self , request ) : <EOL> request = self . _before_record_request ( request ) <EOL> return request and request in self and self . record_mode != '<STR_LIT:all>' and self . rewound <EOL> def play_response ( self , request ) : <EOL> """<STR_LIT>""" <EOL> for index , response in self . _responses ( request ) : <EOL> if self . play_counts [ index ] == <NUM_LIT:0> : <EOL> self . play_counts [ index ] += <NUM_LIT:1> <EOL> return response <EOL> raise UnhandledHTTPRequestError ( <EOL> "<STR_LIT>" <EOL> % ( self . _path , request ) <EOL> ) <EOL> def responses_of ( self , request ) : <EOL> """<STR_LIT>""" <EOL> responses = [ response for index , response in self . _responses ( request ) ] <EOL> if responses : <EOL> return responses <EOL> raise UnhandledHTTPRequestError ( <EOL> "<STR_LIT>" <EOL> % ( self . _path , request ) <EOL> ) <EOL> def _as_dict ( self ) : <EOL> return { "<STR_LIT>" : self . requests , "<STR_LIT>" : self . responses } <EOL> def _save ( self , force = False ) : <EOL> if force or self . dirty : <EOL> save_cassette ( <EOL> self . _path , <EOL> self . _as_dict ( ) , <EOL> serializer = self . _serializer <EOL> ) <EOL> self . dirty = False <EOL> def _load ( self ) : <EOL> try : <EOL> requests , responses = load_cassette ( <EOL> self . _path , <EOL> serializer = self . _serializer <EOL> ) <EOL> for request , response in zip ( requests , responses ) : <EOL> self . append ( request , response ) <EOL> self . dirty = False <EOL> self . rewound = True <EOL> except IOError : <EOL> pass <EOL> def __str__ ( self ) : <EOL> return "<STR_LIT>" . format ( <EOL> len ( self ) <EOL> ) <EOL> def __len__ ( self ) : <EOL> """<STR_LIT>""" <EOL> return len ( self . data ) <EOL> def __contains__ ( self , request ) : <EOL> """<STR_LIT>""" <EOL> for index , response in self . _responses ( request ) : <EOL> if self . play_counts [ index ] == <NUM_LIT:0> : <EOL> return True <EOL> return False </s>
