diff --git a/pytype/overlays/enum_overlay.py b/pytype/overlays/enum_overlay.py index 2baf8350a..7cd990da0 100644 --- a/pytype/overlays/enum_overlay.py +++ b/pytype/overlays/enum_overlay.py @@ -49,7 +49,6 @@ "FlagBoundary", "verify", "property", - "member", "nonmember", "global_enum", "show_flag_values", diff --git a/pytype/stubs/stdlib/enum.pytd b/pytype/stubs/stdlib/enum.pytd index 358ce8c94..03e524ba7 100644 --- a/pytype/stubs/stdlib/enum.pytd +++ b/pytype/stubs/stdlib/enum.pytd @@ -71,7 +71,7 @@ class EnumCheck: ... class FlagBoundary: ... def verify(): ... def property(): ... -def member(): ... -def nonmember(): ... +def member(value: _T) -> _T: ... +def nonmember(value): ... def global_enum(): ... def show_flag_values(): ... diff --git a/pytype/tests/test_enums.py b/pytype/tests/test_enums.py index e5975878a..f6d714774 100644 --- a/pytype/tests/test_enums.py +++ b/pytype/tests/test_enums.py @@ -1634,6 +1634,33 @@ def _missing_(cls, value: object) -> E: ... assert_type(foo.E("FOO"), foo.E) """) + def test_member(self): + """Tests that enum.member acts as a no-op.""" + self.Check(""" + import enum + + class E(enum.Enum): + X = enum.member(1) + + assert_type(E.X, E) + assert_type(E.X.value, int) + """) + + def test_member_with_partial(self): + """Tests that enum.member acts as a no-op even with a partial.""" + self.Check(""" + import enum + import functools + + def foo(x: int) -> int: + return x + 1 + + class E(enum.Enum): + X = enum.member(functools.partial(foo)) + + assert_type(E.X, E) + """) + if __name__ == "__main__": test_base.main()