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