Repository: quora/pyanalyze
Pull request #682 (2023-09-14 23:10:18)

Add EnumName


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 39ebf07..a1d7a61 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,9 @@
 
 ## Unreleased
 
+- Add a `pyanalyze.extensions.EnumName` predicate and infer it
+  as the value for the `.name` attribute on enums. Also fix
+  type inference for enum "properties" on Python 3.11 and up. (#682)
 - Allow `pyanalyze.runtime.is_compatible` to be used to narrow
   types (#681)
 - Fix usage of `assert_type()` with `Any` and with unions of
diff --git a/pyanalyze/__init__.py b/pyanalyze/__init__.py
index b40d599..341ff8e 100644
--- a/pyanalyze/__init__.py
+++ b/pyanalyze/__init__.py
@@ -56,6 +56,7 @@ used(extensions.is_positional)
 used(extensions.is_of_type)
 used(extensions.show_error)
 used(extensions.has_extra_keys)
+used(extensions.EnumName)
 used(value.UNRESOLVED_VALUE)  # keeping it around for now just in case
 used(reexport)
 used(patma)
diff --git a/pyanalyze/annotated_types.py b/pyanalyze/annotated_types.py
index dc4411d..59e7863 100644
--- a/pyanalyze/annotated_types.py
+++ b/pyanalyze/annotated_types.py
@@ -5,7 +5,8 @@ Support for annotations from the annotated_types library.
 """
 from dataclasses import dataclass
 from datetime import datetime, timezone
-from typing import Any, Callable, Iterable, Optional, Union
+import enum
+from typing import Any, Callable, Iterable, Optional, Union, Type
 
 from pyanalyze.value import CanAssign, CanAssignContext, Value, flatten_values
 
@@ -118,6 +119,19 @@ class AnnotatedTypesCheck(CustomCheck):
         )
 
 
+@dataclass(frozen=True)
+class EnumName(AnnotatedTypesCheck):
+    enum_cls: Type[enum.Enum]
+
+    def predicate(self, value: str) -> bool:
+        try:
+            self.enum_cls[value]
+        except KeyError:
+            return False
+        else:
+            return True
+
+
 @dataclass(frozen=True)
 class Gt(AnnotatedTypesCheck):
     value: Any
diff --git a/pyanalyze/attributes.py b/pyanalyze/attributes.py
index 292ddb5..29d7c8d 100644
--- a/pyanalyze/attributes.py
+++ b/pyanalyze/attributes.py
@@ -14,6 +14,7 @@ from typing import Any, Callable, ClassVar, Optional, Sequence, Tuple, Union
 import asynq
 import qcore
 
+from .annotated_types import EnumName
 from .annotations import Context, type_from_annotations, type_from_runtime
 from .options import Options, PyObjectSequenceOption
 from .safe import safe_isinstance, safe_issubclass
@@ -24,6 +25,7 @@ from .value import (
     AnySource,
     AnyValue,
     CallableValue,
+    CustomCheckExtension,
     GenericBases,
     GenericValue,
     HasAttrExtension,
@@ -31,6 +33,7 @@ from .value import (
     KnownValueWithTypeVars,
     MultiValuedValue,
     SyntheticModuleValue,
+    annotate_value,
     set_self,
     SubclassValue,
     TypedValue,
@@ -294,6 +297,8 @@ def _get_attribute_from_typed(
         result = _unwrap_value_from_typed(result, typ, ctx)
     ctx.record_usage(typ, result)
     result = set_self(result, ctx.root_value)
+    if ctx.attr == "name" and safe_issubclass(typ, Enum) and result == TypedValue(str):
+        return annotate_value(result, [CustomCheckExtension(EnumName(typ))])
     return result
 
 
diff --git a/pyanalyze/extensions.py b/pyanalyze/extensions.py
index b0034cc..d593898 100644
--- a/pyanalyze/extensions.py
+++ b/pyanalyze/extensions.py
@@ -12,6 +12,7 @@ import typing
 from collections import defaultdict
 from contextlib import contextmanager
 from dataclasses import dataclass, field
+import enum
 from typing import (
     Any,
     Callable,
@@ -25,12 +26,13 @@ from typing import (
     Sequence,
     Tuple,
     TYPE_CHECKING,
+    Type,
     TypeVar,
     Union,
 )
 
 import typing_extensions
-from typing_extensions import Literal, NoReturn
+from typing_extensions import Literal, NoReturn, Annotated
 
 import pyanalyze
 
@@ -610,3 +612,19 @@ def has_extra_keys(value_type: object = Any) -> Callable[[_T], _T]:
         return cls
 
     return decorator
+
+
+class _EnumName:
+    """A type representing the names of members of an enum.
+
+    Equivalent to a Literal type, but using this will produce nicer error messages
+    for users.
+
+    """
+
+    # TODO after dropping 3.8: switch to a single class with __class_getitem__
+    def __getitem__(self, enum_cls: Type[enum.Enum]) -> Any:
+        return Annotated[str, pyanalyze.annotated_types.EnumName(enum_cls)]
+
+
+EnumName = _EnumName()
diff --git a/pyanalyze/typeshed.py b/pyanalyze/typeshed.py
index 511a3ed..b2355ec 100644
--- a/pyanalyze/typeshed.py
+++ b/pyanalyze/typeshed.py
@@ -7,6 +7,7 @@ Code for getting annotations from typeshed (and from third-party stubs generally
 import ast
 import builtins
 import collections.abc
+import enum
 import inspect
 import sys
 
@@ -15,6 +16,7 @@ from collections.abc import Collection, MutableMapping, Set as AbstractSet
 from dataclasses import dataclass, field, replace
 from enum import Enum, EnumMeta
 from types import GeneratorType, MethodDescriptorType, ModuleType
+import types
 from typing import (
     Any,
     Callable,
@@ -85,6 +87,11 @@ from .value import (
     Value,
 )
 
+PROPERTY_LIKE = {KnownValue(property), KnownValue(types.DynamicClassAttribute)}
+
+if sys.version_info >= (3, 11):
+    PROPERTY_LIKE.add(KnownValue(enum.property))
+
 
 T_co = TypeVar("T_co", covariant=True)
 
@@ -520,7 +527,7 @@ class TypeshedFinder:
             decorators = [
                 self._parse_expr(decorator, mod) for decorator in node.decorator_list
             ]
-            if node.returns and decorators == [KnownValue(property)]:
+            if node.returns and set(decorators) & PROPERTY_LIKE:
                 return self._parse_type(node.returns, mod)
             sig = self._get_signature_from_func_def(
                 node, None, mod, autobind=not on_class


Test patch:
diff --git a/pyanalyze/test_enum.py b/pyanalyze/test_enum.py
index d5295f0..8938c83 100644
--- a/pyanalyze/test_enum.py
+++ b/pyanalyze/test_enum.py
@@ -112,3 +112,24 @@ class TestNarrowing(TestNameCheckVisitorBase):
                 assert_is_value(x, KnownValue(X.a))
                 return
             assert_is_value(x, AnyValue(AnySource.unannotated))
+
+
+class TestEnumName(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test(self):
+        from pyanalyze.extensions import EnumName
+        import enum
+        from typing_extensions import assert_type
+
+        class Rodent(enum.IntEnum):
+            capybara = 1
+            agouti = 2
+
+        def capybara(x: EnumName[Rodent]):
+            pass
+
+        def caller(r: Rodent, s: str):
+            capybara(s)  # E: incompatible_argument
+            capybara(r)  # E: incompatible_argument
+            assert_type(r.name, str)
+            capybara(r.name)