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