Skip to content

Commit

Permalink
feat: support json print in loginfo (#1693)
Browse files Browse the repository at this point in the history
  • Loading branch information
cristianmtr committed Jan 15, 2021
1 parent fa47adb commit c35296e
Showing 1 changed file with 11 additions and 2 deletions.
13 changes: 11 additions & 2 deletions jina/drivers/control.py
Expand Up @@ -3,6 +3,8 @@

import time

from google.protobuf.json_format import MessageToJson

from . import BaseDriver
from .querylang.queryset.dunderkey import dunder_get
from ..excepts import UnknownControlCommand, RuntimeTerminated, NoExplicitMessage
Expand All @@ -22,17 +24,24 @@ def envelope(self) -> 'jina_pb2.EnvelopeProto':
class LogInfoDriver(BaseControlDriver):
"""Log output the request info"""

def __init__(self, key: str = 'request', *args, **kwargs):
def __init__(self, key: str = 'request', json: bool = True, *args, **kwargs):
"""
:param key: (str) that represents a first level or nested key in the dict
:param args:
:param kwargs:
"""
super().__init__(*args, **kwargs)
self.key = key
self.json = json

def __call__(self, *args, **kwargs):
self.logger.info(dunder_get(self.msg.as_pb_object, self.key))
data = dunder_get(self.msg.as_pb_object, self.key)
if self.json:
self.logger.info(
MessageToJson(data)
)
else:
self.logger.info(data)


class WaitDriver(BaseControlDriver):
Expand Down

0 comments on commit c35296e

Please sign in to comment.