From 23aeeb04b5456cb6389a2e96fffd45716ef3190f Mon Sep 17 00:00:00 2001 From: Danny Weinberg Date: Tue, 5 Jan 2021 13:44:51 -0800 Subject: [PATCH 1/7] Add a class attribute hook to the plugin system This adds a hook to modify attributes on *classes* (as opposed to the existing attribute hook, which is for attributes on *instances*). This also adds a test to demonstrate the new behavior. The modifications and added tests were modeled off of https://github.com/python/mypy/commit/3acbf3fe78a61c19ff96754233ada453472004c4. This is intended to solve https://github.com/python/mypy/issues/9645 --- mypy/checkmember.py | 7 ++++++ mypy/plugin.py | 29 ++++++++++++++++++++--- test-data/unit/check-custom-plugin.test | 12 ++++++++++ test-data/unit/plugins/class_attr_hook.py | 20 ++++++++++++++++ 4 files changed, 65 insertions(+), 3 deletions(-) create mode 100644 test-data/unit/plugins/class_attr_hook.py diff --git a/mypy/checkmember.py b/mypy/checkmember.py index 55e0df47ecf9..dd0dd2b7f8de 100644 --- a/mypy/checkmember.py +++ b/mypy/checkmember.py @@ -794,6 +794,13 @@ def analyze_class_attribute_access(itype: Instance, if not mx.is_lvalue: result = analyze_descriptor_access(mx.original_type, result, mx.builtin_type, mx.msg, mx.context, chk=mx.chk) + + # Call the class attribute hook before returning. + fullname = '{}.{}'.format(info.fullname, name) + hook = mx.chk.plugin.get_class_attribute_hook(fullname) + if hook: + result = hook(AttributeContext(get_proper_type(mx.original_type), + result, mx.context, mx.chk)) return result elif isinstance(node.node, Var): mx.not_ready_callback(name, mx.context) diff --git a/mypy/plugin.py b/mypy/plugin.py index 2c39a4548a32..a5b7447dc285 100644 --- a/mypy/plugin.py +++ b/mypy/plugin.py @@ -623,10 +623,10 @@ def get_method_hook(self, fullname: str def get_attribute_hook(self, fullname: str ) -> Optional[Callable[[AttributeContext], Type]]: - """Adjust type of a class attribute. + """Adjust type of an instance attribute. - This method is called with attribute full name using the class where the attribute was - defined (or Var.info.fullname for generated attributes). + This method is called with attribute full name using the class of the instance where + the attribute was 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. @@ -653,6 +653,25 @@ class Derived(Base): """ return None + def get_class_attribute_hook(self, fullname: str + ) -> Optional[Callable[[AttributeContext], Type]]: + """ + 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). + + For example: + + class Cls: + x: Any + + Cls.x + + get_class_attribute_hook is called with '__main__.Cls.x'. + """ + return None + def get_class_decorator_hook(self, fullname: str ) -> Optional[Callable[[ClassDefContext], None]]: """Update class definition for given class decorators. @@ -774,6 +793,10 @@ def get_attribute_hook(self, fullname: str ) -> Optional[Callable[[AttributeContext], Type]]: return self._find_hook(lambda plugin: plugin.get_attribute_hook(fullname)) + def get_class_attribute_hook(self, fullname: str + ) -> Optional[Callable[[AttributeContext], Type]]: + return self._find_hook(lambda plugin: plugin.get_class_attribute_hook(fullname)) + def get_class_decorator_hook(self, fullname: str ) -> Optional[Callable[[ClassDefContext], None]]: return self._find_hook(lambda plugin: plugin.get_class_decorator_hook(fullname)) diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test index 7c85881363d6..db99ffd61ac1 100644 --- a/test-data/unit/check-custom-plugin.test +++ b/test-data/unit/check-custom-plugin.test @@ -803,3 +803,15 @@ reveal_type(f()) # N: Revealed type is "builtins.str" [file mypy.ini] \[mypy] plugins=/test-data/unit/plugins/method_in_decorator.py + +[case testClassAttrPluginFile] +# flags: --config-file tmp/mypy.ini + +class Cls: + attr = 'test' + +reveal_type(Cls.attr) # N: Revealed type is 'builtins.int' +reveal_type(Cls().attr) # N: Revealed type is 'builtins.str' +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/class_attr_hook.py diff --git a/test-data/unit/plugins/class_attr_hook.py b/test-data/unit/plugins/class_attr_hook.py new file mode 100644 index 000000000000..851e5db54977 --- /dev/null +++ b/test-data/unit/plugins/class_attr_hook.py @@ -0,0 +1,20 @@ +from typing import Callable, Optional, Type as TypingType + +from mypy.plugin import AttributeContext, Plugin +from mypy.types import Type as MypyType + + +class ClassAttrPlugin(Plugin): + def get_class_attribute_hook(self, fullname: str + ) -> Optional[Callable[[AttributeContext], MypyType]]: + if fullname == '__main__.Cls.attr': + return my_hook + return None + + +def my_hook(ctx: AttributeContext) -> MypyType: + return ctx.api.named_generic_type('builtins.int', []) + + +def plugin(_version: str) -> TypingType[Plugin]: + return ClassAttrPlugin From f1852e1d347bf277ce701f18a10c63f7b5b69217 Mon Sep 17 00:00:00 2001 From: Danny Weinberg Date: Tue, 5 Jan 2021 16:43:38 -0800 Subject: [PATCH 2/7] Don't import `typing.Type` at runtime to make mypyc happy --- test-data/unit/plugins/class_attr_hook.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/test-data/unit/plugins/class_attr_hook.py b/test-data/unit/plugins/class_attr_hook.py index 851e5db54977..26036d9896e8 100644 --- a/test-data/unit/plugins/class_attr_hook.py +++ b/test-data/unit/plugins/class_attr_hook.py @@ -1,8 +1,12 @@ -from typing import Callable, Optional, Type as TypingType +from typing import Callable, Optional, TYPE_CHECKING from mypy.plugin import AttributeContext, Plugin from mypy.types import Type as MypyType +if TYPE_CHECKING: + # For some reason this fails under the mypyc tests + from typing import Type as TypingType + class ClassAttrPlugin(Plugin): def get_class_attribute_hook(self, fullname: str @@ -16,5 +20,5 @@ def my_hook(ctx: AttributeContext) -> MypyType: return ctx.api.named_generic_type('builtins.int', []) -def plugin(_version: str) -> TypingType[Plugin]: +def plugin(_version: str) -> 'TypingType[Plugin]': return ClassAttrPlugin From bbac561dca45dfda4e7963c3c0969f5ac9351931 Mon Sep 17 00:00:00 2001 From: Danny Weinberg Date: Tue, 5 Jan 2021 18:00:34 -0800 Subject: [PATCH 3/7] mypyc angy, no types 4 u --- test-data/unit/plugins/class_attr_hook.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/test-data/unit/plugins/class_attr_hook.py b/test-data/unit/plugins/class_attr_hook.py index 26036d9896e8..348e5df0ee03 100644 --- a/test-data/unit/plugins/class_attr_hook.py +++ b/test-data/unit/plugins/class_attr_hook.py @@ -1,12 +1,8 @@ -from typing import Callable, Optional, TYPE_CHECKING +from typing import Callable, Optional from mypy.plugin import AttributeContext, Plugin from mypy.types import Type as MypyType -if TYPE_CHECKING: - # For some reason this fails under the mypyc tests - from typing import Type as TypingType - class ClassAttrPlugin(Plugin): def get_class_attribute_hook(self, fullname: str @@ -20,5 +16,5 @@ def my_hook(ctx: AttributeContext) -> MypyType: return ctx.api.named_generic_type('builtins.int', []) -def plugin(_version: str) -> 'TypingType[Plugin]': +def plugin(_version: str): return ClassAttrPlugin From 379aa601b99eef2895420c751eb906f8b232dc5d Mon Sep 17 00:00:00 2001 From: Danny Weinberg Date: Wed, 6 Jan 2021 12:14:52 -0800 Subject: [PATCH 4/7] Handle more exit cases, add more test cases. Some open questions remain, marked with 'TODO RIGHT NOW NOMERGE' --- mypy/checkmember.py | 45 +++++++++----- test-data/unit/check-custom-plugin.test | 81 ++++++++++++++++++++++++- 2 files changed, 109 insertions(+), 17 deletions(-) diff --git a/mypy/checkmember.py b/mypy/checkmember.py index dd0dd2b7f8de..9e192e6e03df 100644 --- a/mypy/checkmember.py +++ b/mypy/checkmember.py @@ -707,10 +707,13 @@ def analyze_class_attribute_access(itype: Instance, if override_info: info = override_info + fullname = '{}.{}'.format(info.fullname, name) + hook = mx.chk.plugin.get_class_attribute_hook(fullname) + node = info.get(name) if not node: if info.fallback_to_any: - return AnyType(TypeOfAny.special_form) + return apply_class_attr_hook(mx, hook, AnyType(TypeOfAny.special_form)) return None is_decorated = isinstance(node.node, Decorator) @@ -735,14 +738,16 @@ def analyze_class_attribute_access(itype: Instance, if info.is_enum and not (mx.is_lvalue or is_decorated or is_method): enum_class_attribute_type = analyze_enum_class_attribute_access(itype, name, mx) if enum_class_attribute_type: - return enum_class_attribute_type + return apply_class_attr_hook(mx, hook, enum_class_attribute_type) t = node.type if t: if isinstance(t, PartialType): symnode = node.node assert isinstance(symnode, Var) - return mx.chk.handle_partial_var_type(t, mx.is_lvalue, symnode, mx.context) + return apply_class_attr_hook(mx, hook, + mx.chk.handle_partial_var_type(t, mx.is_lvalue, symnode, + mx.context)) # Find the class where method/variable was defined. if isinstance(node.node, Decorator): @@ -795,38 +800,38 @@ def analyze_class_attribute_access(itype: Instance, result = analyze_descriptor_access(mx.original_type, result, mx.builtin_type, mx.msg, mx.context, chk=mx.chk) - # Call the class attribute hook before returning. - fullname = '{}.{}'.format(info.fullname, name) - hook = mx.chk.plugin.get_class_attribute_hook(fullname) - if hook: - result = hook(AttributeContext(get_proper_type(mx.original_type), - result, mx.context, mx.chk)) - return result + return apply_class_attr_hook(mx, hook, result) elif isinstance(node.node, Var): + # TODO RIGHT NOW NOMERGE - is it okay to not modify this? mx.not_ready_callback(name, mx.context) return AnyType(TypeOfAny.special_form) if isinstance(node.node, TypeVarExpr): + # TODO RIGHT NOW NOMERGE - is it okay to not modify this? mx.msg.fail(message_registry.CANNOT_USE_TYPEVAR_AS_EXPRESSION.format( info.name, name), mx.context) return AnyType(TypeOfAny.from_error) if isinstance(node.node, TypeInfo): - return type_object_type(node.node, mx.builtin_type) + assert False, "TODO RIGHT NOW NOMERGE - No tests hit this, how to trigger?" + return apply_class_attr_hook(mx, hook, type_object_type(node.node, mx.builtin_type)) if isinstance(node.node, MypyFile): # Reference to a module object. - return mx.builtin_type('types.ModuleType') + assert False, "TODO RIGHT NOW NOMERGE - No tests hit this, how to trigger?" + return apply_class_attr_hook(mx, hook, mx.builtin_type('types.ModuleType')) if (isinstance(node.node, TypeAlias) and isinstance(get_proper_type(node.node.target), Instance)): - return instance_alias_type(node.node, mx.builtin_type) + return apply_class_attr_hook(mx, hook, instance_alias_type(node.node, mx.builtin_type)) if is_decorated: assert isinstance(node.node, Decorator) if node.node.type: - return node.node.type + assert False, "TODO RIGHT NOW NOMERGE - No tests hit this, how to trigger?" + return apply_class_attr_hook(mx, hook, node.node.type) else: + # TODO RIGHT NOW NOMERGE - is it okay to not modify this? mx.not_ready_callback(name, mx.context) return AnyType(TypeOfAny.from_error) else: @@ -837,7 +842,17 @@ def analyze_class_attribute_access(itype: Instance, # unannotated implicit class methods we do this here. if node.node.is_class: typ = bind_self(typ, is_classmethod=True) - return typ + return apply_class_attr_hook(mx, hook, typ) + + +def apply_class_attr_hook(mx: MemberContext, + hook: Optional[Callable[[AttributeContext], Type]], + result: Type, + ) -> Optional[Type]: + if hook: + result = hook(AttributeContext(get_proper_type(mx.original_type), + result, mx.context, mx.chk)) + return result def analyze_enum_class_attribute_access(itype: Instance, diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test index db99ffd61ac1..5be976dec478 100644 --- a/test-data/unit/check-custom-plugin.test +++ b/test-data/unit/check-custom-plugin.test @@ -804,14 +804,91 @@ reveal_type(f()) # N: Revealed type is "builtins.str" \[mypy] plugins=/test-data/unit/plugins/method_in_decorator.py -[case testClassAttrPluginFile] +[case testClassAttrPluginClassVar] # flags: --config-file tmp/mypy.ini +from typing import Type + class Cls: attr = 'test' + unchanged = 'test' -reveal_type(Cls.attr) # N: Revealed type is 'builtins.int' reveal_type(Cls().attr) # N: Revealed type is 'builtins.str' +reveal_type(Cls.attr) # N: Revealed type is 'builtins.int' +reveal_type(Cls.unchanged) # N: Revealed type is 'builtins.str' +x: Type[Cls] +reveal_type(x.attr) # N: Revealed type is 'builtins.int' +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/class_attr_hook.py + +[case testClassAttrPluginMethod] +# flags: --config-file tmp/mypy.ini + +class Cls: + def attr(self) -> None: + pass + +reveal_type(Cls.attr) # N: Revealed type is 'builtins.int' +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/class_attr_hook.py + +[case testClassAttrPluginEnum] +# flags: --config-file tmp/mypy.ini + +import enum + +class Cls(enum.Enum): + attr = 'test' + +reveal_type(Cls.attr) # N: Revealed type is 'builtins.int' +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/class_attr_hook.py + +[case testClassAttrPluginMetaclass] +# flags: --config-file tmp/mypy.ini + +from typing import Any, Type +class M(type): + attr = 'test' + +B: Any +class Cls(B, metaclass=M): + pass + +reveal_type(Cls.attr) # N: Revealed type is 'builtins.int' +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/class_attr_hook.py + +[case testClassAttrPluginPartialType] +# flags: --config-file tmp/mypy.ini + +class Cls: + attr = None + def f(self) -> int: + return Cls.attr + 1 + +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/class_attr_hook.py + +[case testClassAttrPluginAliasRef] +# flags: --config-file tmp/mypy.ini + +from typing import Generic, TypeVar, Type + +T = TypeVar('T') +class G(Generic[T]): + pass + +class Cls: + attr = G[int] + +x: Type[Cls] +reveal_type(x.attr) # N: Revealed type is 'builtins.int' [file mypy.ini] \[mypy] plugins=/test-data/unit/plugins/class_attr_hook.py From 987207accedca9bdffbadb9858d99d12146145c6 Mon Sep 17 00:00:00 2001 From: Danny Weinberg Date: Fri, 15 Jan 2021 09:12:07 -0800 Subject: [PATCH 5/7] Resolve open questions, mainly moving toward not using the hook for them. --- mypy/checkmember.py | 12 +++--------- test-data/unit/check-custom-plugin.test | 18 ------------------ 2 files changed, 3 insertions(+), 27 deletions(-) diff --git a/mypy/checkmember.py b/mypy/checkmember.py index 9e192e6e03df..7e7cb5b48ce7 100644 --- a/mypy/checkmember.py +++ b/mypy/checkmember.py @@ -802,36 +802,30 @@ def analyze_class_attribute_access(itype: Instance, return apply_class_attr_hook(mx, hook, result) elif isinstance(node.node, Var): - # TODO RIGHT NOW NOMERGE - is it okay to not modify this? mx.not_ready_callback(name, mx.context) return AnyType(TypeOfAny.special_form) if isinstance(node.node, TypeVarExpr): - # TODO RIGHT NOW NOMERGE - is it okay to not modify this? mx.msg.fail(message_registry.CANNOT_USE_TYPEVAR_AS_EXPRESSION.format( info.name, name), mx.context) return AnyType(TypeOfAny.from_error) if isinstance(node.node, TypeInfo): - assert False, "TODO RIGHT NOW NOMERGE - No tests hit this, how to trigger?" - return apply_class_attr_hook(mx, hook, type_object_type(node.node, mx.builtin_type)) + return type_object_type(node.node, mx.builtin_type) if isinstance(node.node, MypyFile): # Reference to a module object. - assert False, "TODO RIGHT NOW NOMERGE - No tests hit this, how to trigger?" - return apply_class_attr_hook(mx, hook, mx.builtin_type('types.ModuleType')) + return mx.builtin_type('types.ModuleType') if (isinstance(node.node, TypeAlias) and isinstance(get_proper_type(node.node.target), Instance)): - return apply_class_attr_hook(mx, hook, instance_alias_type(node.node, mx.builtin_type)) + return instance_alias_type(node.node, mx.builtin_type) if is_decorated: assert isinstance(node.node, Decorator) if node.node.type: - assert False, "TODO RIGHT NOW NOMERGE - No tests hit this, how to trigger?" return apply_class_attr_hook(mx, hook, node.node.type) else: - # TODO RIGHT NOW NOMERGE - is it okay to not modify this? mx.not_ready_callback(name, mx.context) return AnyType(TypeOfAny.from_error) else: diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test index 5be976dec478..bddaab482171 100644 --- a/test-data/unit/check-custom-plugin.test +++ b/test-data/unit/check-custom-plugin.test @@ -874,21 +874,3 @@ class Cls: [file mypy.ini] \[mypy] plugins=/test-data/unit/plugins/class_attr_hook.py - -[case testClassAttrPluginAliasRef] -# flags: --config-file tmp/mypy.ini - -from typing import Generic, TypeVar, Type - -T = TypeVar('T') -class G(Generic[T]): - pass - -class Cls: - attr = G[int] - -x: Type[Cls] -reveal_type(x.attr) # N: Revealed type is 'builtins.int' -[file mypy.ini] -\[mypy] -plugins=/test-data/unit/plugins/class_attr_hook.py From a976739d5c6c8b4d5ee7e724bb290dcc64354fdd Mon Sep 17 00:00:00 2001 From: Danny Weinberg Date: Mon, 5 Apr 2021 18:29:05 -0700 Subject: [PATCH 6/7] Update quotation marks for rebase --- test-data/unit/check-custom-plugin.test | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test index bddaab482171..01e9d44dc81a 100644 --- a/test-data/unit/check-custom-plugin.test +++ b/test-data/unit/check-custom-plugin.test @@ -813,11 +813,11 @@ class Cls: attr = 'test' unchanged = 'test' -reveal_type(Cls().attr) # N: Revealed type is 'builtins.str' -reveal_type(Cls.attr) # N: Revealed type is 'builtins.int' -reveal_type(Cls.unchanged) # N: Revealed type is 'builtins.str' +reveal_type(Cls().attr) # N: Revealed type is "builtins.str" +reveal_type(Cls.attr) # N: Revealed type is "builtins.int" +reveal_type(Cls.unchanged) # N: Revealed type is "builtins.str" x: Type[Cls] -reveal_type(x.attr) # N: Revealed type is 'builtins.int' +reveal_type(x.attr) # N: Revealed type is "builtins.int" [file mypy.ini] \[mypy] plugins=/test-data/unit/plugins/class_attr_hook.py @@ -829,7 +829,7 @@ class Cls: def attr(self) -> None: pass -reveal_type(Cls.attr) # N: Revealed type is 'builtins.int' +reveal_type(Cls.attr) # N: Revealed type is "builtins.int" [file mypy.ini] \[mypy] plugins=/test-data/unit/plugins/class_attr_hook.py @@ -842,7 +842,7 @@ import enum class Cls(enum.Enum): attr = 'test' -reveal_type(Cls.attr) # N: Revealed type is 'builtins.int' +reveal_type(Cls.attr) # N: Revealed type is "builtins.int" [file mypy.ini] \[mypy] plugins=/test-data/unit/plugins/class_attr_hook.py @@ -858,7 +858,7 @@ B: Any class Cls(B, metaclass=M): pass -reveal_type(Cls.attr) # N: Revealed type is 'builtins.int' +reveal_type(Cls.attr) # N: Revealed type is "builtins.int" [file mypy.ini] \[mypy] plugins=/test-data/unit/plugins/class_attr_hook.py From 1816cf6a0377b34fb27550f1be9eb3b7a726111f Mon Sep 17 00:00:00 2001 From: Danny Weinberg Date: Tue, 22 Feb 2022 19:32:27 +0000 Subject: [PATCH 7/7] Add documentation and additional test --- docs/source/extending_mypy.rst | 4 ++++ mypy/plugin.py | 2 +- test-data/unit/check-custom-plugin.test | 20 +++++++++++++++++++- 3 files changed, 24 insertions(+), 2 deletions(-) diff --git a/docs/source/extending_mypy.rst b/docs/source/extending_mypy.rst index 5c59bef506cc..00c328be7728 100644 --- a/docs/source/extending_mypy.rst +++ b/docs/source/extending_mypy.rst @@ -198,6 +198,10 @@ fields which already exist on the class. *Exception:* if :py:meth:`__getattr__ < :py:meth:`__getattribute__ ` is a method on the class, the hook is called for all fields which do not refer to methods. +**get_class_attribute_hook()** is similar to above, but for attributes on classes rather than instances. +Unlike above, this does not have special casing for :py:meth:`__getattr__ ` or +:py:meth:`__getattribute__ `. + **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 to match runtime behaviour: diff --git a/mypy/plugin.py b/mypy/plugin.py index 7f3a5270d85f..8a4f39186085 100644 --- a/mypy/plugin.py +++ b/mypy/plugin.py @@ -677,7 +677,7 @@ class Cls: Cls.x - get_class_attribute_hook is called with '__main__.Cls.x'. + get_class_attribute_hook is called with '__main__.Cls.x' as fullname. """ return None diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test index c9eda99322a2..6f8dac77c442 100644 --- a/test-data/unit/check-custom-plugin.test +++ b/test-data/unit/check-custom-plugin.test @@ -946,7 +946,7 @@ reveal_type(Cls.attr) # N: Revealed type is "builtins.int" \[mypy] plugins=/test-data/unit/plugins/class_attr_hook.py -[case testClassAttrPluginMetaclass] +[case testClassAttrPluginMetaclassAnyBase] # flags: --config-file tmp/mypy.ini from typing import Any, Type @@ -962,6 +962,24 @@ reveal_type(Cls.attr) # N: Revealed type is "builtins.int" \[mypy] plugins=/test-data/unit/plugins/class_attr_hook.py +[case testClassAttrPluginMetaclassRegularBase] +# flags: --config-file tmp/mypy.ini + +from typing import Any, Type +class M(type): + attr = 'test' + +class B: + attr = None + +class Cls(B, metaclass=M): + pass + +reveal_type(Cls.attr) # N: Revealed type is "builtins.int" +[file mypy.ini] +\[mypy] +plugins=/test-data/unit/plugins/class_attr_hook.py + [case testClassAttrPluginPartialType] # flags: --config-file tmp/mypy.ini