Skip to content

Instantly share code, notes, and snippets.

@hzhu212
Last active December 11, 2020 08:21
Show Gist options
  • Save hzhu212/795db2eb13b4151fee9c078d2a2e6919 to your computer and use it in GitHub Desktop.
Save hzhu212/795db2eb13b4151fee9c078d2a2e6919 to your computer and use it in GitHub Desktop.
python protobuf wrapper

A wrapper for Python protobuf object:

class PbMessage(object):
    """protobuf 消息的包装类,封装一些通用方法"""

    def __init__(self, pb_class):
        self._pb_class = pb_class
        self._pb_obj = None
        self._raw_msg = None
        self._msg = None

    def _deserialize(self):
        if not self._pb_obj:
            self._pb_obj = self._pb_class()
        self._pb_obj.ParseFromString(self._raw_msg)
        self._msg = self._pb_obj

    def update(self, msg_body):
        self._raw_msg = msg_body
        self._deserialize()
        return self
    
    def to_dict(self):
        from google.protobuf.json_format import MessageToDict
        return MessageToDict(self._msg, preserving_proto_field_name=True)

    def _get(self, key):
        try:
            return self._msg[key]
        except (TypeError, KeyError):
            return getattr(self._msg, key)

    def get(self, key, default=None):
        try:
            return self._get(key)
        except (TypeError, KeyError, AttributeError):
            return default
            
    def __getattr__(self, key):
        try:
            return self._get(key)
        except (TypeError, KeyError, AttributeError):
            raise AttributeError(f'{self.__class__.__name__} object has no attribute named {key!r}')

    def __getitem__(self, key):
        try:
            return self._get(key)
        except (TypeError, KeyError, AttributeError):
            raise KeyError(f'{self.__class__.__name__} object has no item indexed by key {key!r}')

    def __str__(self):
        return f'<{self.__class__.__name__}(msg_body={self._raw_msg}, pb_class={self._pb_class})>'
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment