Skip to content

Commit 9bfeea8

Browse files
OpenHTF Ownerscopybara-github
authored andcommitted
Prevent calling as_base_types on class objects in openhtf.util.data.
Update openhtf.util.data to check that an object is not a class before calling its as_base_types method. This ensures that class objects themselves are converted to their string representations instead of attempting to call the method on the class. Also add corresponding unit tests. PiperOrigin-RevId: 963749692
1 parent adef239 commit 9bfeea8

2 files changed

Lines changed: 14 additions & 2 deletions

File tree

openhtf/util/data.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -163,7 +163,7 @@ def convert_to_base_types(obj,
163163
# Because it's *really* annoying to pass a single string accidentally.
164164
assert not isinstance(ignore_keys, str), 'Pass a real iterable!'
165165

166-
if hasattr(obj, 'as_base_types'):
166+
if hasattr(obj, 'as_base_types') and not inspect.isclass(obj):
167167
return obj.as_base_types()
168168
if hasattr(obj, '_asdict') and not inspect.isclass(obj):
169169
obj = obj._asdict()

test/util/data_test.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,11 @@ def __init__(self):
4646
def as_base_types(self):
4747
return self.value
4848

49+
class ClassWithAsBaseTypes(object):
50+
51+
def as_base_types(self):
52+
return 'called on instance'
53+
4954
@attr.s(slots=True, frozen=True)
5055
class FrozenAttr(object):
5156
value = attr.ib(type=int)
@@ -72,7 +77,8 @@ class EnumClass(enum.Enum):
7277
'special': SpecialBaseTypes('must_not_be_present'),
7378
'not_copied': not_copied,
7479
'enum': EnumClass.A,
75-
80+
'class_with_as_base_types_instance': ClassWithAsBaseTypes(),
81+
'class_with_as_base_types_class': ClassWithAsBaseTypes,
7682
# Some plugs such as UserInputPlug will return None as a response to
7783
# AsDict().
7884
'none_dict': AsDict(),
@@ -99,6 +105,12 @@ class StrEnumClass(enum.StrEnum):
99105
self.assertIsInstance(converted['special'], dict)
100106
self.assertEqual(converted['special'], {'safe_value': True})
101107
self.assertIs(converted['not_copied'], not_copied.value)
108+
self.assertEqual(
109+
converted['class_with_as_base_types_instance'], 'called on instance'
110+
)
111+
self.assertEqual(
112+
converted['class_with_as_base_types_class'], str(ClassWithAsBaseTypes)
113+
)
102114

103115
self.assertIsNone(converted['none_dict'])
104116

0 commit comments

Comments
 (0)