Skip to content

Commit 35e019e

Browse files
authored
feat: add custom __dir__ for messages and message classes (#289)
During development, it can be convenient to inspect objects and types directly to determine what methods and attributes they have using the dir() builtin command in a debugger or a REPL. Because proto-plus messages wrap their fields using __getattr__, the proto fields are not visible by default and must be explicitly exposed to dir().
1 parent 28aa3b2 commit 35e019e

File tree

2 files changed

+113
-0
lines changed

2 files changed

+113
-0
lines changed

proto/message.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,27 @@ def __prepare__(mcls, name, bases, **kwargs):
273273
def meta(cls):
274274
return cls._meta
275275

276+
def __dir__(self):
277+
names = set(dir(type))
278+
names.update(
279+
(
280+
"meta",
281+
"pb",
282+
"wrap",
283+
"serialize",
284+
"deserialize",
285+
"to_json",
286+
"from_json",
287+
"to_dict",
288+
"copy_from",
289+
)
290+
)
291+
desc = self.pb().DESCRIPTOR
292+
names.update(t.name for t in desc.nested_types)
293+
names.update(e.name for e in desc.enum_types)
294+
295+
return names
296+
276297
def pb(cls, obj=None, *, coerce: bool = False):
277298
"""Return the underlying protobuf Message class or instance.
278299
@@ -520,6 +541,29 @@ def __init__(
520541
# Create the internal protocol buffer.
521542
super().__setattr__("_pb", self._meta.pb(**params))
522543

544+
def __dir__(self):
545+
desc = type(self).pb().DESCRIPTOR
546+
names = {f_name for f_name in self._meta.fields.keys()}
547+
names.update(m.name for m in desc.nested_types)
548+
names.update(e.name for e in desc.enum_types)
549+
names.update(dir(object()))
550+
# Can't think of a better way of determining
551+
# the special methods than manually listing them.
552+
names.update(
553+
(
554+
"__bool__",
555+
"__contains__",
556+
"__dict__",
557+
"__getattr__",
558+
"__getstate__",
559+
"__module__",
560+
"__setstate__",
561+
"__weakref__",
562+
)
563+
)
564+
565+
return names
566+
523567
def __bool__(self):
524568
"""Return True if any field is truthy, False otherwise."""
525569
return any(k in self and getattr(self, k) for k in self._meta.fields.keys())

tests/test_message.py

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -346,3 +346,72 @@ class Squid(proto.Message):
346346

347347
with pytest.raises(TypeError):
348348
Mollusc.Squid.copy_from(m.squid, (("mass_kg", 20)))
349+
350+
351+
def test_dir():
352+
class Mollusc(proto.Message):
353+
class Class(proto.Enum):
354+
UNKNOWN = 0
355+
GASTROPOD = 1
356+
BIVALVE = 2
357+
CEPHALOPOD = 3
358+
359+
class Arm(proto.Message):
360+
length_cm = proto.Field(proto.INT32, number=1)
361+
362+
mass_kg = proto.Field(proto.INT32, number=1)
363+
class_ = proto.Field(Class, number=2)
364+
arms = proto.RepeatedField(Arm, number=3)
365+
366+
expected = (
367+
{
368+
# Fields and nested message and enum types
369+
"arms",
370+
"class_",
371+
"mass_kg",
372+
"Arm",
373+
"Class",
374+
}
375+
| {
376+
# Other methods and attributes
377+
"__bool__",
378+
"__contains__",
379+
"__dict__",
380+
"__getattr__",
381+
"__getstate__",
382+
"__module__",
383+
"__setstate__",
384+
"__weakref__",
385+
}
386+
| set(dir(object))
387+
) # Gets the long tail of dunder methods and attributes.
388+
389+
actual = set(dir(Mollusc()))
390+
391+
# Check instance names
392+
assert actual == expected
393+
394+
# Check type names
395+
expected = (
396+
set(dir(type))
397+
| {
398+
# Class methods from the MessageMeta metaclass
399+
"copy_from",
400+
"deserialize",
401+
"from_json",
402+
"meta",
403+
"pb",
404+
"serialize",
405+
"to_dict",
406+
"to_json",
407+
"wrap",
408+
}
409+
| {
410+
# Nested message and enum types
411+
"Arm",
412+
"Class",
413+
}
414+
)
415+
416+
actual = set(dir(Mollusc))
417+
assert actual == expected

0 commit comments

Comments
 (0)