1# Copyright 2019 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 15from typing import Type 16import enum 17import warnings 18 19 20class EnumRule: 21 """A marshal for converting between integer values and enum values.""" 22 23 def __init__(self, enum_class: Type[enum.IntEnum]): 24 self._enum = enum_class 25 26 def to_python(self, value, *, absent: bool = None): 27 if isinstance(value, int) and not isinstance(value, self._enum): 28 try: 29 # Coerce the int on the wire to the enum value. 30 return self._enum(value) 31 except ValueError: 32 # Since it is possible to add values to enums, we do 33 # not want to flatly error on this. 34 # 35 # However, it is useful to make some noise about it so 36 # the user realizes that an unexpected value came along. 37 warnings.warn( 38 "Unrecognized {name} enum value: {value}".format( 39 name=self._enum.__name__, value=value, 40 ) 41 ) 42 return value 43 44 def to_proto(self, value): 45 # Accept enum values and coerce to the pure integer. 46 # This is not strictly necessary (protocol buffers can take these 47 # objects as they subclass int) but nevertheless seems like the 48 # right thing to do. 49 if isinstance(value, self._enum): 50 return value.value 51 52 # If a string is provided that matches an enum value, coerce it 53 # to the enum value. 54 if isinstance(value, str): 55 return self._enum[value].value 56 57 # We got a pure integer; pass it on. 58 return value 59