codekingpro/portable-devtools

1# Copyright 2018 Google LLC2#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 at6#7#     https://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# 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 and13# limitations under the License.14 15import abc16import enum17 18from google.protobuf import message19from google.protobuf import duration_pb220from google.protobuf import timestamp_pb221from google.protobuf import field_mask_pb222from google.protobuf import struct_pb223from google.protobuf import wrappers_pb224 25from proto.marshal import compat26from proto.marshal.collections import MapComposite27from proto.marshal.collections import Repeated28from proto.marshal.collections import RepeatedComposite29 30from proto.marshal.rules import bytes as pb_bytes31from proto.marshal.rules import stringy_numbers32from proto.marshal.rules import dates33from proto.marshal.rules import struct34from proto.marshal.rules import wrappers35from proto.marshal.rules import field_mask36from proto.primitives import ProtoType37 38 39class Rule(abc.ABC):40    """Abstract class definition for marshal rules."""41 42    @classmethod43    def __subclasshook__(cls, C):44        if hasattr(C, "to_python") and hasattr(C, "to_proto"):45            return True46        return NotImplemented47 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 essentially54    translates between these: it keeps a registry of common protocol buffers55    and their Python representations, and translates back and forth.56 57    The protocol buffer class is always the "key" in this relationship; when58    presenting a message, the declared field types are used to determine59    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 should65    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) and77        a ``rule``; an object with a ``to_python`` and ``to_proto`` method.78        Each method should return the appropriate Python or protocol buffer79        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 zero88        arguments.89 90        Args:91            proto_type (type): A protocol buffer message type.92            rule: A marshal object93        """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] = rule105            return106 107        # Create an inner function that will register an instance of the108        # 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_class120 121        return register_rule_class122 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 they148        # create RepeatedComposite and MapComposite instances directly and149        # 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/decode155        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/2679159        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/349171        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                    break176        return rule177 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 isinstance182        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        if value_type in compat.map_composite_types:192            return MapComposite(value, marshal=self)193        return self.get_rule(proto_type=proto_type).to_python(value, absent=absent)194 195    def to_proto(self, proto_type, value, *, strict: bool = False):196        # The protos in google/protobuf/struct.proto are exceptional cases,197        # because they can and should represent themselves as lists and dicts.198        # These cases are handled in their rule classes.199        if proto_type not in (200            struct_pb2.Value,201            struct_pb2.ListValue,202            struct_pb2.Struct,203        ):204            # For our repeated and map view objects, simply return the205            # underlying pb.206            if isinstance(value, (Repeated, MapComposite)):207                return value.pb208 209            # Convert lists and tuples recursively.210            if isinstance(value, (list, tuple)):211                return type(value)(self.to_proto(proto_type, i) for i in value)212 213        # Convert dictionaries recursively when the proto type is a map.214        # This is slightly more complicated than converting a list or tuple215        # because we have to step through the magic that protocol buffers does.216        #217        # Essentially, a type of map<string, Foo> will show up here as218        # a FoosEntry with a `key` field, `value` field, and a `map_entry`219        # annotation. We need to do the conversion based on the `value`220        # field's type.221        if isinstance(value, dict) and (222            proto_type.DESCRIPTOR.has_options223            and proto_type.DESCRIPTOR.GetOptions().map_entry224        ):225            recursive_type = type(proto_type().value)226            return {k: self.to_proto(recursive_type, v) for k, v in value.items()}227 228        pb_value = self.get_rule(proto_type=proto_type).to_proto(value)229 230        # Sanity check: If we are in strict mode, did we get the value we want?231        if strict and not isinstance(pb_value, proto_type):232            raise TypeError(233                "Parameter must be instance of the same class; "234                "expected {expected}, got {got}".format(235                    expected=proto_type.__name__,236                    got=pb_value.__class__.__name__,237                ),238            )239        # Return the final value.240        return pb_value241 242 243class Marshal(BaseMarshal):244    """The translator between protocol buffer and Python instances.245 246    The bulk of the implementation is in :class:`BaseMarshal`. This class247    adds identity tracking: multiple instantiations of :class:`Marshal` with248    the same name will provide the same instance.249    """250 251    _instances = {}252 253    def __new__(cls, *, name: str):254        """Create a marshal instance.255 256        Args:257            name (str): The name of the marshal. Instantiating multiple258                marshals with the same ``name`` argument will provide the259                same marshal each time.260        """261        klass = cls._instances.get(name)262        if klass is None:263            klass = cls._instances[name] = super().__new__(cls)264 265        return klass266 267    def __init__(self, *, name: str):268        """Instantiate a marshal.269 270        Args:271            name (str): The name of the marshal. Instantiating multiple272                marshals with the same ``name`` argument will provide the273                same marshal each time.274        """275        self._name = name276        if not hasattr(self, "_rules"):277            super().__init__()278 279 280class NoopRule:281    """A catch-all rule that does nothing."""282 283    def to_python(self, pb_value, *, absent: bool = None):284        return pb_value285 286    def to_proto(self, value):287        return value288 289 290__all__ = ("Marshal",)291