From 89d9e32340414989d142bb0d8c74a30135c08467 Mon Sep 17 00:00:00 2001 From: tsudol Date: Thu, 30 Sep 2021 17:51:59 -0700 Subject: [PATCH] Allow enum members in Literals. We still allow unsolvables, because enum members will continue to show up as those until the overlay is enabled. PiperOrigin-RevId: 400070986 --- pytype/overlays/typing_overlay.py | 4 +++- pytype/tests/test_typing2.py | 9 +++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/pytype/overlays/typing_overlay.py b/pytype/overlays/typing_overlay.py index f80a259e1..03589e85b 100644 --- a/pytype/overlays/typing_overlay.py +++ b/pytype/overlays/typing_overlay.py @@ -687,7 +687,7 @@ def _build_value(self, node, inner, ellipses): values = [] errors = [] for i, param in enumerate(inner): - # TODO(b/173742489): Once pytype has proper support for enums, we should + # TODO(b/173742489): Once the enum overlay is enabled, we should # stop allowing unsolvable and handle enums here. if (param == self.vm.convert.none or isinstance(param, abstract.LiteralClass) or @@ -696,6 +696,8 @@ def _build_value(self, node, inner, ellipses): elif (isinstance(param, abstract.ConcreteValue) and isinstance(param.pyval, (int, str, bytes))): value = abstract.LiteralClass(param, self.vm) + elif isinstance(param, abstract.Instance) and param.cls.is_enum: + value = abstract.LiteralClass(param, self.vm) else: if i in ellipses: invalid_param = "..." diff --git a/pytype/tests/test_typing2.py b/pytype/tests/test_typing2.py index 05defffcf..dc6fe2976 100644 --- a/pytype/tests/test_typing2.py +++ b/pytype/tests/test_typing2.py @@ -900,6 +900,15 @@ def f(x: Literal["x", "y"]): f(x) """) + def test_enum(self): + # Requires the enum overlay + self.Check(""" + import enum + from typing import Literal + class M(enum.Enum): + A = 1 + x: Literal[M.A] + """) if __name__ == "__main__": test_base.main()