def test_enums():
L = descriptor_pb2.SourceCodeInfo.Location
enum_pb = descriptor_pb2.EnumDescriptorProto(
name='Silly',
value=(
descriptor_pb2.EnumValueDescriptorProto(name='ZERO', number=0),
descriptor_pb2.EnumValueDescriptorProto(name='ONE', number=1),
descriptor_pb2.EnumValueDescriptorProto(name='THREE', number=3),
))
fdp = make_file_pb2(package='google.enum.v1',
enums=(enum_pb, ),
locations=(
L(path=(5, 0),
leading_comments='This is the Silly enum.'),
L(path=(5, 0, 2, 0),
leading_comments='This is the zero value.'),
L(path=(5, 0, 2, 1),
leading_comments='This is the one value.'),
))
proto = api.Proto.build(fdp, file_to_generate=True)
assert len(proto.enums) == 1
enum = proto.enums['google.enum.v1.Silly']
assert enum.meta.doc == 'This is the Silly enum.'
assert isinstance(enum, wrappers.EnumType)
assert len(enum.values) == 3
assert all([isinstance(i, wrappers.EnumValueType) for i in enum.values])
assert enum.values[0].name == 'ZERO'
assert enum.values[0].meta.doc == 'This is the zero value.'
assert enum.values[1].name == 'ONE'
assert enum.values[1].meta.doc == 'This is the one value.'
assert enum.values[2].name == 'THREE'
assert enum.values[2].meta.doc == ''
def get_enum() -> wrappers.EnumType:
enum_value_pbs = [
descriptor_pb2.EnumValueDescriptorProto(name='RED', number=1),
descriptor_pb2.EnumValueDescriptorProto(name='GREEN', number=2),
descriptor_pb2.EnumValueDescriptorProto(name='BLUE', number=3),
]
enum_pb = descriptor_pb2.EnumDescriptorProto(
name='Color',
value=enum_value_pbs,
)
return wrappers.EnumType(
enum_pb=enum_pb,
values=[wrappers.EnumValueType(enum_value_pb=evpb)
for evpb in enum_value_pbs],
)
def make_enum(
name: str,
package: str = 'foo.bar.v1',
module: str = 'baz',
values: typing.Tuple[str, int] = (),
meta: metadata.Metadata = None,
) -> wrappers.EnumType:
enum_value_pbs = [
desc.EnumValueDescriptorProto(name=i[0], number=i[1]) for i in values
]
enum_pb = desc.EnumDescriptorProto(
name=name,
value=enum_value_pbs,
)
return wrappers.EnumType(
enum_pb=enum_pb,
values=[
wrappers.EnumValueType(enum_value_pb=evpb)
for evpb in enum_value_pbs
],
meta=meta or metadata.Metadata(address=metadata.Address(
name=name,
package=tuple(package.split('.')),
module=module,
)),
)
def test_mock_value_enum():
values = [
descriptor_pb2.EnumValueDescriptorProto(name='UNSPECIFIED', number=0),
descriptor_pb2.EnumValueDescriptorProto(name='SPECIFIED', number=1),
]
enum = wrappers.EnumType(
values=[wrappers.EnumValueType(enum_value_pb=i) for i in values],
enum_pb=descriptor_pb2.EnumDescriptorProto(value=values),
meta=metadata.Metadata(address=metadata.Address(
module='bogus',
name='Enumerable',
)),
)
field = make_field(
type='TYPE_ENUM',
type_name='bogus.Enumerable',
enum=enum,
)
assert field.mock_value == 'bogus.Enumerable.SPECIFIED'
def make_enum_pb2(name: str, *values: typing.Sequence[str],
**kwargs) -> desc.EnumDescriptorProto:
enum_value_pbs = [
desc.EnumValueDescriptorProto(name=n, number=i)
for i, n in enumerate(values)
]
enum_pb = desc.EnumDescriptorProto(name=name,
value=enum_value_pbs,
**kwargs)
return enum_pb
def enum_factory(name: str, variants: Iterable[str]) -> wrappers.EnumType:
enum_pb = descriptor_pb2.EnumDescriptorProto(
name=name,
value=tuple(
descriptor_pb2.EnumValueDescriptorProto(name=v, number=i)
for i, v in enumerate(variants)))
enum = wrappers.EnumType(enum_pb=enum_pb,
values=[
wrappers.EnumValueType(enum_value_pb=v)
for v in enum_pb.value
])
return enum
def __new__(mcls, name, bases, attrs):
# Do not do any special behavior for `proto.Enum` itself.
if bases[0] == enum.IntEnum:
return super().__new__(mcls, name, bases, attrs)
# Get the essential information about the proto package, and where
# this component belongs within the file.
package, marshal = _package_info.compile(name, attrs)
# Determine the local path of this proto component within the file.
local_path = tuple(attrs.get("__qualname__", name).split("."))
# Sanity check: We get the wrong full name if a class is declared
# inside a function local scope; correct this.
if "<locals>" in local_path:
ix = local_path.index("<locals>")
local_path = local_path[: ix - 1] + local_path[ix + 1 :]
# Determine the full name in protocol buffers.
full_name = ".".join((package,) + local_path).lstrip(".")
filename = _file_info._FileInfo.proto_file_name(
attrs.get("__module__", name.lower())
)
# Retrieve any enum options.
# We expect something that looks like an EnumOptions message,
# either an actual instance or a dict-like representation.
pb_options = "_pb_options"
opts = attrs.pop(pb_options, {})
# This is the only portable way to remove the _pb_options name
# from the enum attrs.
# In 3.7 onwards, we can define an _ignore_ attribute and do some
# mucking around with that.
if pb_options in attrs._member_names:
idx = attrs._member_names.index(pb_options)
attrs._member_names.pop(idx)
# Make the descriptor.
enum_desc = descriptor_pb2.EnumDescriptorProto(
name=name,
# Note: the superclass ctor removes the variants, so get them now.
# Note: proto3 requires that the first variant value be zero.
value=sorted(
(
descriptor_pb2.EnumValueDescriptorProto(name=name, number=number)
# Minor hack to get all the enum variants out.
for name, number in attrs.items()
if isinstance(number, int)
),
key=lambda v: v.number,
),
options=opts,
)
file_info = _file_info._FileInfo.maybe_add_descriptor(filename, package)
if len(local_path) == 1:
file_info.descriptor.enum_type.add().MergeFrom(enum_desc)
else:
file_info.nested_enum[local_path] = enum_desc
# Run the superclass constructor.
cls = super().__new__(mcls, name, bases, attrs)
# We can't just add a "_meta" element to attrs because the Enum
# machinery doesn't know what to do with a non-int value.
# The pb is set later, in generate_file_pb
cls._meta = _EnumInfo(full_name=full_name, pb=None)
file_info.enums[full_name] = cls
# Register the enum with the marshal.
marshal.register(cls, EnumRule(cls))
# Generate the descriptor for the file if it is ready.
if file_info.ready(new_class=cls):
file_info.generate_file_pb(new_class=cls, fallback_salt=full_name)
# Done; return the class.
return cls
def make_enum_value_pb2(name: str, number: int,
**kwargs) -> desc.EnumValueDescriptorProto:
"""Mock an EnumValueDescriptorProto that is used to create wrappers.EnumValue"""
return desc.EnumValueDescriptorProto(name=name, number=number)
