"""Helper functions for clients in Jina."""fromfunctoolsimportwrapsfrominspectimportsignaturefromtypingimportCallable,Optionalimportwarningsfromjina.exceptsimportBadClientCallbackfromjina.logging.loggerimportJinaLoggerfromjina.protoimportjina_pb2fromjina.types.request.dataimportResponsefromjina.helperimportget_rich_console
[docs]defpprint_routes(resp:'Response',stack_limit:int=3):"""Pretty print routes with :mod:`prettytable`, fallback to :func:`print`. :param resp: the :class:`Response` object :param stack_limit: traceback limit """routes=resp.routesfromrich.tableimportTablefromrichimportboxtable=Table(box=box.SIMPLE)forvin('Executor','Time','Exception'):table.add_column(v)forrouteinroutes:status_icon='🟢'ifroute.status.code==jina_pb2.StatusProto.ERROR:status_icon='🔴'elifroute.status.code==jina_pb2.StatusProto.ERROR_CHAINED:status_icon='⚪'table.add_row(f'{status_icon}{route.executor}',f'{route.start_time.ToMilliseconds()-routes[0].start_time.ToMilliseconds()}ms',''.join(route.status.exception.stacks[-stack_limit:]),)console=get_rich_console()console.print(table)
def_safe_callback(func:Callable,continue_on_error:bool,logger)->Callable:@wraps(func)def_arg_wrapper(*args,**kwargs):try:returnfunc(*args,**kwargs)exceptExceptionasex:err_msg=f'uncaught exception in callback {func.__name__}(): {ex!r}'ifcontinue_on_error:logger.error(err_msg)else:raiseBadClientCallback(err_msg)fromexreturn_arg_wrapper
[docs]defcallback_exec(response,on_done:Callable,on_error:Callable,on_always:Callable,continue_on_error:bool,logger:JinaLogger,)->None:"""Execute the callback with the response. :param response: the response :param on_done: the on_done callback :param on_error: the on_error callback :param on_always: the on_always callback :param continue_on_error: whether to continue on error :param logger: a logger instance """ifon_errorandresponse.header.status.code>=jina_pb2.StatusProto.ERROR:@wraps(on_error)defon_error_wrap(resp):on_error(resp,None)_safe_callback(on_error_wrap,continue_on_error,logger)(response)elifon_doneandresponse.header.status.code==jina_pb2.StatusProto.SUCCESS:_safe_callback(on_done,continue_on_error,logger)(response)ifon_always:_safe_callback(on_always,continue_on_error,logger)(response)
[docs]defcallback_exec_on_error(on_error:Callable,exception:Exception,logger:JinaLogger,response:Optional=None,)->None:"""Execute the on_error callback with the response, Use when an error outside the response status was thrown. :param on_error: the on_error callback :param exception: the exception with was thrown and led to the call of on_error :param logger: a logger instance :param response: the response """@wraps(on_error)defon_error_wrap(resp):on_error(resp,exception)_safe_callback(on_error_wrap,False,logger)(response)