60 lines
2.2 KiB
Python
60 lines
2.2 KiB
Python
# Copyright 2019 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from typing import Type
|
|
import enum
|
|
import warnings
|
|
|
|
|
|
class EnumRule:
|
|
"""A marshal for converting between integer values and enum values."""
|
|
|
|
def __init__(self, enum_class: Type[enum.IntEnum]):
|
|
self._enum = enum_class
|
|
|
|
def to_python(self, value, *, absent: bool = None):
|
|
if isinstance(value, int) and not isinstance(value, self._enum):
|
|
try:
|
|
# Coerce the int on the wire to the enum value.
|
|
return self._enum(value)
|
|
except ValueError:
|
|
# Since it is possible to add values to enums, we do
|
|
# not want to flatly error on this.
|
|
#
|
|
# However, it is useful to make some noise about it so
|
|
# the user realizes that an unexpected value came along.
|
|
warnings.warn(
|
|
"Unrecognized {name} enum value: {value}".format(
|
|
name=self._enum.__name__,
|
|
value=value,
|
|
)
|
|
)
|
|
return value
|
|
|
|
def to_proto(self, value):
|
|
# Accept enum values and coerce to the pure integer.
|
|
# This is not strictly necessary (protocol buffers can take these
|
|
# objects as they subclass int) but nevertheless seems like the
|
|
# right thing to do.
|
|
if isinstance(value, self._enum):
|
|
return value.value
|
|
|
|
# If a string is provided that matches an enum value, coerce it
|
|
# to the enum value.
|
|
if isinstance(value, str):
|
|
return self._enum[value].value
|
|
|
|
# We got a pure integer; pass it on.
|
|
return value
|