Skip to content

Commit

Permalink
fix: adding enums to a repeated field does not raise a TypeError (#202)
Browse files Browse the repository at this point in the history
Fixes issue #201, where enums added to a repeated field triggered a
TypeError because they were coverted to integers during marshaling.
  • Loading branch information
software-dov authored Mar 4, 2021
1 parent 03653da commit 2a10bbe
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 1 deletion.
2 changes: 1 addition & 1 deletion proto/marshal/collections/repeated.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,5 +174,5 @@ def __setitem__(self, key, value):

def insert(self, index: int, value):
"""Insert ``value`` in the sequence before ``index``."""
pb_value = self._marshal.to_proto(self._pb_type, value, strict=True)
pb_value = self._marshal.to_proto(self._pb_type, value)
self.pb.insert(index, pb_value)
24 changes: 24 additions & 0 deletions tests/test_marshal_strict.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
# Copyright (C) 2021 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
#
# http://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.

import proto
from proto.marshal.marshal import BaseMarshal
import pytest


def test_strict_to_proto():
m = BaseMarshal()

with pytest.raises(TypeError):
m.to_proto(dict, None, strict=True)
32 changes: 32 additions & 0 deletions tests/test_marshal_types_enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,3 +58,35 @@ class Foo(proto.Enum):
with mock.patch.object(warnings, "warn") as warn:
assert enum_rule.to_python(4) == 4
warn.assert_called_once_with("Unrecognized Foo enum value: 4")


def test_enum_append():
class Bivalve(proto.Enum):
CLAM = 0
OYSTER = 1

class MolluscContainer(proto.Message):
bivalves = proto.RepeatedField(proto.ENUM, number=1, enum=Bivalve,)

mc = MolluscContainer()
clam = Bivalve.CLAM
mc.bivalves.append(clam)
mc.bivalves.append(1)

assert mc.bivalves == [clam, Bivalve.OYSTER]


def test_enum_map_insert():
class Bivalve(proto.Enum):
CLAM = 0
OYSTER = 1

class MolluscContainer(proto.Message):
bivalves = proto.MapField(proto.STRING, proto.ENUM, number=1, enum=Bivalve,)

mc = MolluscContainer()
clam = Bivalve.CLAM
mc.bivalves["clam"] = clam
mc.bivalves["oyster"] = 1

assert mc.bivalves == {"clam": clam, "oyster": Bivalve.OYSTER}

0 comments on commit 2a10bbe

Please sign in to comment.