Skip to content

Commit

Permalink
feat: add custom __dir__ for messages and message classes (#289)
Browse files Browse the repository at this point in the history
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().
  • Loading branch information
software-dov authored Feb 7, 2022
1 parent 28aa3b2 commit 35e019e
Show file tree
Hide file tree
Showing 2 changed files with 113 additions and 0 deletions.
44 changes: 44 additions & 0 deletions proto/message.py
Original file line number Diff line number Diff line change
Expand Up @@ -273,6 +273,27 @@ def __prepare__(mcls, name, bases, **kwargs):
def meta(cls):
return cls._meta

def __dir__(self):
names = set(dir(type))
names.update(
(
"meta",
"pb",
"wrap",
"serialize",
"deserialize",
"to_json",
"from_json",
"to_dict",
"copy_from",
)
)
desc = self.pb().DESCRIPTOR
names.update(t.name for t in desc.nested_types)
names.update(e.name for e in desc.enum_types)

return names

def pb(cls, obj=None, *, coerce: bool = False):
"""Return the underlying protobuf Message class or instance.
Expand Down Expand Up @@ -520,6 +541,29 @@ def __init__(
# Create the internal protocol buffer.
super().__setattr__("_pb", self._meta.pb(**params))

def __dir__(self):
desc = type(self).pb().DESCRIPTOR
names = {f_name for f_name in self._meta.fields.keys()}
names.update(m.name for m in desc.nested_types)
names.update(e.name for e in desc.enum_types)
names.update(dir(object()))
# Can't think of a better way of determining
# the special methods than manually listing them.
names.update(
(
"__bool__",
"__contains__",
"__dict__",
"__getattr__",
"__getstate__",
"__module__",
"__setstate__",
"__weakref__",
)
)

return names

def __bool__(self):
"""Return True if any field is truthy, False otherwise."""
return any(k in self and getattr(self, k) for k in self._meta.fields.keys())
Expand Down
69 changes: 69 additions & 0 deletions tests/test_message.py
Original file line number Diff line number Diff line change
Expand Up @@ -346,3 +346,72 @@ class Squid(proto.Message):

with pytest.raises(TypeError):
Mollusc.Squid.copy_from(m.squid, (("mass_kg", 20)))


def test_dir():
class Mollusc(proto.Message):
class Class(proto.Enum):
UNKNOWN = 0
GASTROPOD = 1
BIVALVE = 2
CEPHALOPOD = 3

class Arm(proto.Message):
length_cm = proto.Field(proto.INT32, number=1)

mass_kg = proto.Field(proto.INT32, number=1)
class_ = proto.Field(Class, number=2)
arms = proto.RepeatedField(Arm, number=3)

expected = (
{
# Fields and nested message and enum types
"arms",
"class_",
"mass_kg",
"Arm",
"Class",
}
| {
# Other methods and attributes
"__bool__",
"__contains__",
"__dict__",
"__getattr__",
"__getstate__",
"__module__",
"__setstate__",
"__weakref__",
}
| set(dir(object))
) # Gets the long tail of dunder methods and attributes.

actual = set(dir(Mollusc()))

# Check instance names
assert actual == expected

# Check type names
expected = (
set(dir(type))
| {
# Class methods from the MessageMeta metaclass
"copy_from",
"deserialize",
"from_json",
"meta",
"pb",
"serialize",
"to_dict",
"to_json",
"wrap",
}
| {
# Nested message and enum types
"Arm",
"Class",
}
)

actual = set(dir(Mollusc))
assert actual == expected

0 comments on commit 35e019e

Please sign in to comment.