Use the get_attribute hook for dynamic classes (#6371)

lincolnq · ilevkivskyi · commit 2a32d7382616 · 2019-02-10T17:20:08.000Z
Whenever we find `__getattr__` or `__getattribute__` on an instance, the getattribute hook is called to give the hook a chance to modify the return type of one of those functions. Closes #6259, #5910
diff --git a/docs/source/extending_mypy.rst b/docs/source/extending_mypy.rst
@@ -184,11 +184,11 @@ For example in this code:
 mypy will call ``get_method_signature_hook("ctypes.Array.__setitem__")``
 so that the plugin can mimic the ``ctypes`` auto-convert behavior.
 
-**get_attribute_hook** can be used to give more precise type of an instance
-attribute. Note however, that this method is only called for variables that
-already exist in the class symbol table. If you want to add some generated
-variables/methods to the symbol table you can use one of the three hooks
-below.
+**get_attribute_hook** overrides instance member field lookups and property
+access (not assignments, and not method calls). This hook is only called for
+fields which already exist on the class. *Exception:* if ``__getattr__`` or
+``__getattribute__`` is a method on the class, the hook is called for all
+fields which do not refer to methods.
 
 **get_class_decorator_hook()** can be used to update class definition for
 given class decorators. For example, you can add some attributes to the class
diff --git a/mypy/checkmember.py b/mypy/checkmember.py
@@ -340,7 +340,15 @@ def analyze_member_var_access(name: str,
                     typ = map_instance_to_supertype(itype, method.info)
                     getattr_type = expand_type_by_instance(bound_method, typ)
                     if isinstance(getattr_type, CallableType):
-                        return getattr_type.ret_type
+                        result = getattr_type.ret_type
+
+                        # Call the attribute hook before returning.
+                        fullname = '{}.{}'.format(method.info.fullname(), name)
+                        hook = mx.chk.plugin.get_attribute_hook(fullname)
+                        if hook:
+                            result = hook(AttributeContext(mx.original_type, result,
+                                                           mx.context, mx.chk))
+                        return result
         else:
             setattr_meth = info.get_method('__setattr__')
             if setattr_meth and setattr_meth.info.fullname() != 'builtins.object':
diff --git a/mypy/plugin.py b/mypy/plugin.py
@@ -440,8 +440,15 @@ def get_attribute_hook(self, fullname: str
         """Adjust type of a class attribute.
 
         This method is called with attribute full name using the class where the attribute was
-        defined (or Var.info.fullname() for generated attributes). Currently, this hook is only
-        called for names that exist in the class MRO, for example in:
+        defined (or Var.info.fullname() for generated attributes).
+
+        For classes without __getattr__ or __getattribute__, this hook is only called for
+        names of fields/properties (but not methods) that exist in the instance MRO.
+
+        For classes that implement __getattr__ or __getattribute__, this hook is called
+        for all fields/properties, including nonexistent ones (but still not methods).
+
+        For example:
 
             class Base:
                 x: Any
@@ -454,7 +461,9 @@ class Derived(Base):
             var.x
             var.y
 
-        this method is only called with '__main__.Base.x'.
+        get_attribute_hook is called with '__main__.Base.x' and '__main__.Base.y'.
+        However, if we had not implemented __getattr__ on Base, you would only get
+        the callback for 'var.x'; 'var.y' would produce an error without calling the hook.
         """
         return None
 
diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test
@@ -143,19 +143,49 @@ tmp/mypy.ini:2: error: Return value of "plugin" must be a subclass of "mypy.plug
 [case testAttributeTypeHookPlugin]
 # flags: --config-file tmp/mypy.ini
 from typing import Callable
-from m import Signal
+from m import Signal, DerivedSignal
 s: Signal[Callable[[int], None]] = Signal()
 s(1)
 s('') # E: Argument 1 has incompatible type "str"; expected "int"
+
+ds: DerivedSignal[Callable[[int], None]] = DerivedSignal()
+ds('') # E: Argument 1 has incompatible type "str"; expected "int"
 [file m.py]
 from typing import TypeVar, Generic, Callable
 T = TypeVar('T', bound=Callable[..., None])
 class Signal(Generic[T]):
     __call__: Callable[..., None]  # This type is replaced by the plugin
+
+class DerivedSignal(Signal[T]): ...
 [file mypy.ini]
 [[mypy]
 plugins=<ROOT>/test-data/unit/plugins/attrhook.py
 
+[case testAttributeHookPluginForDynamicClass]
+# flags: --config-file tmp/mypy.ini
+from m import Magic, DerivedMagic
+
+magic = Magic()
+reveal_type(magic.magic_field)  # E: Revealed type is 'builtins.str'
+reveal_type(magic.non_magic_method())  # E: Revealed type is 'builtins.int'
+reveal_type(magic.non_magic_field)  # E: Revealed type is 'builtins.int'
+magic.nonexistent_field  # E: Field does not exist
+reveal_type(magic.fallback_example)  # E: Revealed type is 'Any'
+reveal_type(DerivedMagic().magic_field)  # E: Revealed type is 'builtins.str'
+[file m.py]
+from typing import Any
+class Magic:
+    # Triggers plugin infrastructure:
+    def __getattr__(self, x: Any) -> Any: ...
+    def non_magic_method(self) -> int: ...
+    non_magic_field: int
+
+class DerivedMagic(Magic): ...
+
+[file mypy.ini]
+[[mypy]
+plugins=<ROOT>/test-data/unit/plugins/attrhook2.py
+
 [case testTypeAnalyzeHookPlugin]
 # flags: --config-file tmp/mypy.ini
 from typing import Callable
diff --git a/test-data/unit/plugins/attrhook.py b/test-data/unit/plugins/attrhook.py
@@ -12,9 +12,9 @@ def get_attribute_hook(self, fullname: str) -> Optional[Callable[[AttributeConte
 
 
 def signal_call_callback(ctx: AttributeContext) -> Type:
-    if isinstance(ctx.type, Instance) and ctx.type.type.fullname() == 'm.Signal':
+    if isinstance(ctx.type, Instance):
         return ctx.type.args[0]
-    return ctx.inferred_attr_type
+    return ctx.default_attr_type
 
 
 def plugin(version):
diff --git a/test-data/unit/plugins/attrhook2.py b/test-data/unit/plugins/attrhook2.py
@@ -0,0 +1,26 @@
+from typing import Optional, Callable
+
+from mypy.plugin import Plugin, AttributeContext
+from mypy.types import Type, AnyType, TypeOfAny
+
+
+class AttrPlugin(Plugin):
+    def get_attribute_hook(self, fullname: str) -> Optional[Callable[[AttributeContext], Type]]:
+        if fullname == 'm.Magic.magic_field':
+            return magic_field_callback
+        if fullname == 'm.Magic.nonexistent_field':
+            return nonexistent_field_callback
+        return None
+
+
+def magic_field_callback(ctx: AttributeContext) -> Type:
+    return ctx.api.named_generic_type("builtins.str", [])
+
+
+def nonexistent_field_callback(ctx: AttributeContext) -> Type:
+    ctx.api.fail("Field does not exist", ctx.context)
+    return AnyType(TypeOfAny.from_error)
+
+
+def plugin(version):
+    return AttrPlugin
