Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/proto/marshal/rules/enums.py: 42%

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__, 

40 value=value, 

41 ) 

42 ) 

43 return value 

44 

45 def to_proto(self, value): 

46 # Accept enum values and coerce to the pure integer. 

47 # This is not strictly necessary (protocol buffers can take these 

48 # objects as they subclass int) but nevertheless seems like the 

49 # right thing to do. 

50 if isinstance(value, self._enum): 

51 return value.value 

52 

53 # If a string is provided that matches an enum value, coerce it 

54 # to the enum value. 

55 if isinstance(value, str): 

56 return self._enum[value].value 

57 

58 # We got a pure integer; pass it on. 

59 return value