1# Copyright 2018 Google LLC 
    2# 
    3# Licensed under the Apache License, Version 2.0 (the "License"); 
    4# you may not use this file except in compliance with the License. 
    5# You may obtain a copy of the License at 
    6# 
    7#     https://www.apache.org/licenses/LICENSE-2.0 
    8# 
    9# Unless required by applicable law or agreed to in writing, software 
    10# distributed under the License is distributed on an "AS IS" BASIS, 
    11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
    12# See the License for the specific language governing permissions and 
    13# limitations under the License. 
    14 
    15import abc 
    16import enum 
    17 
    18from google.protobuf import message 
    19from google.protobuf import duration_pb2 
    20from google.protobuf import timestamp_pb2 
    21from google.protobuf import field_mask_pb2 
    22from google.protobuf import struct_pb2 
    23from google.protobuf import wrappers_pb2 
    24 
    25from proto.marshal import compat 
    26from proto.marshal.collections import MapComposite 
    27from proto.marshal.collections import Repeated 
    28from proto.marshal.collections import RepeatedComposite 
    29 
    30from proto.marshal.rules import bytes as pb_bytes 
    31from proto.marshal.rules import stringy_numbers 
    32from proto.marshal.rules import dates 
    33from proto.marshal.rules import struct 
    34from proto.marshal.rules import wrappers 
    35from proto.marshal.rules import field_mask 
    36from proto.primitives import ProtoType 
    37 
    38 
    39class Rule(abc.ABC): 
    40    """Abstract class definition for marshal rules.""" 
    41 
    42    @classmethod 
    43    def __subclasshook__(cls, C): 
    44        if hasattr(C, "to_python") and hasattr(C, "to_proto"): 
    45            return True 
    46        return NotImplemented 
    47 
    48 
    49class BaseMarshal: 
    50    """The base class to translate between protobuf and Python classes. 
    51 
    52    Protocol buffers defines many common types (e.g. Timestamp, Duration) 
    53    which also exist in the Python standard library. The marshal essentially 
    54    translates between these: it keeps a registry of common protocol buffers 
    55    and their Python representations, and translates back and forth. 
    56 
    57    The protocol buffer class is always the "key" in this relationship; when 
    58    presenting a message, the declared field types are used to determine 
    59    whether a value should be transformed into another class. Similarly, 
    60    when accepting a Python value (when setting a field, for example), 
    61    the declared field type is still used. This means that, if appropriate, 
    62    multiple protocol buffer types may use the same Python type. 
    63 
    64    The primary implementation of this is :class:`Marshal`, which should 
    65    usually be used instead of this class directly. 
    66    """ 
    67 
    68    def __init__(self): 
    69        self._rules = {} 
    70        self._noop = NoopRule() 
    71        self.reset() 
    72 
    73    def register(self, proto_type: type, rule: Rule = None): 
    74        """Register a rule against the given ``proto_type``. 
    75 
    76        This function expects a ``proto_type`` (the descriptor class) and 
    77        a ``rule``; an object with a ``to_python`` and ``to_proto`` method. 
    78        Each method should return the appropriate Python or protocol buffer 
    79        type, and be idempotent (e.g. accept either type as input). 
    80 
    81        This function can also be used as a decorator:: 
    82 
    83            @marshal.register(timestamp_pb2.Timestamp) 
    84            class TimestampRule: 
    85                ... 
    86 
    87        In this case, the class will be initialized for you with zero 
    88        arguments. 
    89 
    90        Args: 
    91            proto_type (type): A protocol buffer message type. 
    92            rule: A marshal object 
    93        """ 
    94        # If a rule was provided, register it and be done. 
    95        if rule: 
    96            # Ensure the rule implements Rule. 
    97            if not isinstance(rule, Rule): 
    98                raise TypeError( 
    99                    "Marshal rule instances must implement " 
    100                    "`to_proto` and `to_python` methods." 
    101                ) 
    102 
    103            # Register the rule. 
    104            self._rules[proto_type] = rule 
    105            return 
    106 
    107        # Create an inner function that will register an instance of the 
    108        # marshal class to this object's registry, and return it. 
    109        def register_rule_class(rule_class: type): 
    110            # Ensure the rule class is a valid rule. 
    111            if not issubclass(rule_class, Rule): 
    112                raise TypeError( 
    113                    "Marshal rule subclasses must implement " 
    114                    "`to_proto` and `to_python` methods." 
    115                ) 
    116 
    117            # Register the rule class. 
    118            self._rules[proto_type] = rule_class() 
    119            return rule_class 
    120 
    121        return register_rule_class 
    122 
    123    def reset(self): 
    124        """Reset the registry to its initial state.""" 
    125        self._rules.clear() 
    126 
    127        # Register date and time wrappers. 
    128        self.register(timestamp_pb2.Timestamp, dates.TimestampRule()) 
    129        self.register(duration_pb2.Duration, dates.DurationRule()) 
    130 
    131        # Register FieldMask wrappers. 
    132        self.register(field_mask_pb2.FieldMask, field_mask.FieldMaskRule()) 
    133 
    134        # Register nullable primitive wrappers. 
    135        self.register(wrappers_pb2.BoolValue, wrappers.BoolValueRule()) 
    136        self.register(wrappers_pb2.BytesValue, wrappers.BytesValueRule()) 
    137        self.register(wrappers_pb2.DoubleValue, wrappers.DoubleValueRule()) 
    138        self.register(wrappers_pb2.FloatValue, wrappers.FloatValueRule()) 
    139        self.register(wrappers_pb2.Int32Value, wrappers.Int32ValueRule()) 
    140        self.register(wrappers_pb2.Int64Value, wrappers.Int64ValueRule()) 
    141        self.register(wrappers_pb2.StringValue, wrappers.StringValueRule()) 
    142        self.register(wrappers_pb2.UInt32Value, wrappers.UInt32ValueRule()) 
    143        self.register(wrappers_pb2.UInt64Value, wrappers.UInt64ValueRule()) 
    144 
    145        # Register the google.protobuf.Struct wrappers. 
    146        # 
    147        # These are aware of the marshal that created them, because they 
    148        # create RepeatedComposite and MapComposite instances directly and 
    149        # need to pass the marshal to them. 
    150        self.register(struct_pb2.Value, struct.ValueRule(marshal=self)) 
    151        self.register(struct_pb2.ListValue, struct.ListValueRule(marshal=self)) 
    152        self.register(struct_pb2.Struct, struct.StructRule(marshal=self)) 
    153 
    154        # Special case for bytes to allow base64 encode/decode 
    155        self.register(ProtoType.BYTES, pb_bytes.BytesRule()) 
    156 
    157        # Special case for int64 from strings because of dict round trip. 
    158        # See https://github.com/protocolbuffers/protobuf/issues/2679 
    159        for rule_class in stringy_numbers.STRINGY_NUMBER_RULES: 
    160            self.register(rule_class._proto_type, rule_class()) 
    161 
    162    def get_rule(self, proto_type): 
    163        # Rules are needed to convert values between proto-plus and pb. 
    164        # Retrieve the rule for the specified proto type. 
    165        # The NoopRule will be used when a rule is not found. 
    166        rule = self._rules.get(proto_type, self._noop) 
    167 
    168        # If we don't find a rule, also check under `_instances` 
    169        # in case there is a rule in another package. 
    170        # See https://github.com/googleapis/proto-plus-python/issues/349 
    171        if rule == self._noop and hasattr(self, "_instances"): 
    172            for _, instance in self._instances.items(): 
    173                rule = instance._rules.get(proto_type, self._noop) 
    174                if rule != self._noop: 
    175                    break 
    176        return rule 
    177 
    178    def to_python(self, proto_type, value, *, absent: bool = None): 
    179        # Internal protobuf has its own special type for lists of values. 
    180        # Return a view around it that implements MutableSequence. 
    181        value_type = type(value)  # Minor performance boost over isinstance 
    182        if value_type in compat.repeated_composite_types: 
    183            return RepeatedComposite(value, marshal=self) 
    184        if value_type in compat.repeated_scalar_types: 
    185            if isinstance(proto_type, type): 
    186                return RepeatedComposite(value, marshal=self, proto_type=proto_type) 
    187            else: 
    188                return Repeated(value, marshal=self) 
    189 
    190        # Same thing for maps of messages. 
    191        # See https://github.com/protocolbuffers/protobuf/issues/16596 
    192        # We need to look up the name of the type in compat.map_composite_type_names 
    193        # as class `MessageMapContainer` is no longer exposed 
    194        # This is done to avoid taking a breaking change in proto-plus. 
    195        if ( 
    196            value_type in compat.map_composite_types 
    197            or value_type.__name__ in compat.map_composite_type_names 
    198        ): 
    199            return MapComposite(value, marshal=self) 
    200        return self.get_rule(proto_type=proto_type).to_python(value, absent=absent) 
    201 
    202    def to_proto(self, proto_type, value, *, strict: bool = False): 
    203        # The protos in google/protobuf/struct.proto are exceptional cases, 
    204        # because they can and should represent themselves as lists and dicts. 
    205        # These cases are handled in their rule classes. 
    206        if proto_type not in ( 
    207            struct_pb2.Value, 
    208            struct_pb2.ListValue, 
    209            struct_pb2.Struct, 
    210        ): 
    211            # For our repeated and map view objects, simply return the 
    212            # underlying pb. 
    213            if isinstance(value, (Repeated, MapComposite)): 
    214                return value.pb 
    215 
    216            # Convert lists and tuples recursively. 
    217            if isinstance(value, (list, tuple)): 
    218                return type(value)(self.to_proto(proto_type, i) for i in value) 
    219 
    220        # Convert dictionaries recursively when the proto type is a map. 
    221        # This is slightly more complicated than converting a list or tuple 
    222        # because we have to step through the magic that protocol buffers does. 
    223        # 
    224        # Essentially, a type of map<string, Foo> will show up here as 
    225        # a FoosEntry with a `key` field, `value` field, and a `map_entry` 
    226        # annotation. We need to do the conversion based on the `value` 
    227        # field's type. 
    228        if isinstance(value, dict) and ( 
    229            proto_type.DESCRIPTOR.has_options 
    230            and proto_type.DESCRIPTOR.GetOptions().map_entry 
    231        ): 
    232            recursive_type = type(proto_type().value) 
    233            return {k: self.to_proto(recursive_type, v) for k, v in value.items()} 
    234 
    235        pb_value = self.get_rule(proto_type=proto_type).to_proto(value) 
    236 
    237        # Sanity check: If we are in strict mode, did we get the value we want? 
    238        if strict and not isinstance(pb_value, proto_type): 
    239            raise TypeError( 
    240                "Parameter must be instance of the same class; " 
    241                "expected {expected}, got {got}".format( 
    242                    expected=proto_type.__name__, 
    243                    got=pb_value.__class__.__name__, 
    244                ), 
    245            ) 
    246        # Return the final value. 
    247        return pb_value 
    248 
    249 
    250class Marshal(BaseMarshal): 
    251    """The translator between protocol buffer and Python instances. 
    252 
    253    The bulk of the implementation is in :class:`BaseMarshal`. This class 
    254    adds identity tracking: multiple instantiations of :class:`Marshal` with 
    255    the same name will provide the same instance. 
    256    """ 
    257 
    258    _instances = {} 
    259 
    260    def __new__(cls, *, name: str): 
    261        """Create a marshal instance. 
    262 
    263        Args: 
    264            name (str): The name of the marshal. Instantiating multiple 
    265                marshals with the same ``name`` argument will provide the 
    266                same marshal each time. 
    267        """ 
    268        klass = cls._instances.get(name) 
    269        if klass is None: 
    270            klass = cls._instances[name] = super().__new__(cls) 
    271 
    272        return klass 
    273 
    274    def __init__(self, *, name: str): 
    275        """Instantiate a marshal. 
    276 
    277        Args: 
    278            name (str): The name of the marshal. Instantiating multiple 
    279                marshals with the same ``name`` argument will provide the 
    280                same marshal each time. 
    281        """ 
    282        self._name = name 
    283        if not hasattr(self, "_rules"): 
    284            super().__init__() 
    285 
    286 
    287class NoopRule: 
    288    """A catch-all rule that does nothing.""" 
    289 
    290    def to_python(self, pb_value, *, absent: bool = None): 
    291        return pb_value 
    292 
    293    def to_proto(self, value): 
    294        return value 
    295 
    296 
    297__all__ = ("Marshal",)