Repository: quora/pyanalyze
Pull request #485 (2022-03-02 22:15:00)

Infer callable signatures for objects with a `__getattr__` method


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index aacfc6d..062d355 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Infer callable signatures for objects with a
+  `__getattr__` method (#485)
 - Do not treat attributes that raise an exception on access
   as nonexistent (#481)
 - Improve detection of unhashable dict keys and set members (#469)
diff --git a/pyanalyze/arg_spec.py b/pyanalyze/arg_spec.py
index 65e4fe7..2b0f850 100644
--- a/pyanalyze/arg_spec.py
+++ b/pyanalyze/arg_spec.py
@@ -13,6 +13,7 @@ from . import implementation
 from .safe import (
     all_of_type,
     is_newtype,
+    safe_equals,
     safe_hasattr,
     safe_issubclass,
     is_typing_name,
@@ -58,6 +59,7 @@ import asynq
 from collections.abc import Awaitable
 import contextlib
 from dataclasses import dataclass, replace
+import enum
 import qcore
 import inspect
 import sys
@@ -81,6 +83,8 @@ from unittest import mock
 # types.MethodWrapperType in 3.7+
 MethodWrapperType = type(object().__str__)
 
+_ENUM_CALL = enum.Enum.__call__
+
 
 @used  # exposed as an API
 @contextlib.contextmanager
@@ -572,7 +576,7 @@ class ArgSpecCache:
                 if evaluator_sig is not None:
                     return evaluator_sig
 
-        if isinstance(obj, tuple) or hasattr(obj, "__getattr__"):
+        if isinstance(obj, tuple):
             return None  # lost cause
 
         # Cythonized methods, e.g. fn.asynq
@@ -618,6 +622,16 @@ class ArgSpecCache:
                     original_fn, impl, is_asynq, in_overload_resolution
                 )
 
+        # Special case for EnumMeta.__call__. Ideally this should be generalized.
+        if (
+            safe_isinstance(obj, type)
+            and safe_issubclass(obj, enum.Enum)
+            and safe_equals(obj.__call__, _ENUM_CALL)
+        ):
+            return self._cached_get_argspec(
+                _ENUM_CALL, impl, is_asynq, in_overload_resolution
+            )
+
         allow_call = FunctionsSafeToCall.contains(obj, self.options) or (
             safe_isinstance(obj, type) and safe_issubclass(obj, self.safe_bases)
         )


Test patch:
diff --git a/pyanalyze/test_enum.py b/pyanalyze/test_enum.py
new file mode 100644
index 0000000..800172d
--- /dev/null
+++ b/pyanalyze/test_enum.py
@@ -0,0 +1,29 @@
+# static analysis: ignore
+from .implementation import assert_is_value
+from .value import SubclassValue, TypedValue
+from .test_name_check_visitor import TestNameCheckVisitorBase
+from .test_node_visitor import assert_passes
+
+
+class TestEnum(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_functional(self):
+        from enum import Enum
+
+        def capybara():
+            X = Enum("X", ["a", "b", "c"])
+            assert_is_value(X, SubclassValue(TypedValue(Enum)))
+
+    @assert_passes()
+    def test_call(self):
+        from enum import Enum
+
+        class X(Enum):
+            a = 1
+            b = 2
+
+        def capybara():
+            assert_is_value(X(1), TypedValue(X))
+            # This should be an error, but the typeshed
+            # stubs are too lenient.
+            assert_is_value(X(None), TypedValue(X))