Add initial mypy plugin for Record

Author: bryanforbes

asyncpg/mypy/__init__.py

@@ -0,0 +1,49 @@
+import typing
+import mypy.nodes
+import mypy.plugin
+import mypy.types
+
+from . import common
+from . import hooks
+from . import utils
+
+
+class AsyncpgPlugin(mypy.plugin.Plugin):
+    def get_method_hook(self, fullname: str) \
+            -> typing.Optional[common.MethodHook]:
+        class_name, _, method_name = fullname.rpartition('.')
+        symbol = self.lookup_fully_qualified(class_name)
+
+        if symbol and isinstance(symbol.node, mypy.nodes.TypeInfo) and \
+                utils.is_record(symbol.node):
+            if method_name == '__getitem__':
+                return hooks.record_getitem
+            if method_name == 'get':
+                return hooks.record_get
+
+        return None
+
+    def get_attribute_hook(self, fullname: str) \
+            -> typing.Optional[common.AttributeHook]:
+        class_name, _, _ = fullname.rpartition('.')
+        symbol = self.lookup_fully_qualified(class_name)
+
+        if symbol and isinstance(symbol.node, mypy.nodes.TypeInfo) and \
+                symbol.node.has_base(common.RECORD_NAME):
+            return hooks.record_attribute
+
+        return None
+
+    def get_customize_class_mro_hook(self, fullname: str) \
+            -> typing.Optional[common.ClassDefHook]:
+        return hooks.mark_record
+
+    def get_base_class_hook(self, fullname: str) \
+            -> typing.Optional[common.ClassDefHook]:
+        if fullname == common.RECORD_NAME:
+            return hooks.record_final_attributes
+        return None
+
+
+def plugin(version: str) -> typing.Type[mypy.plugin.Plugin]:
+    return AsyncpgPlugin

asyncpg/mypy/common.py

@@ -0,0 +1,25 @@
+import mypy.nodes
+import mypy.plugin
+import mypy.types
+import typing
+import typing_extensions
+
+RECORD_NAME = 'asyncpg.protocol.protocol.Record'  # type: typing_extensions.Final  # noqa: E501
+MethodPairType = typing.Tuple[typing.List[mypy.nodes.Argument],
+                              mypy.types.Type]
+FieldPairType = typing.Tuple[str, mypy.types.Type]
+
+
+class MethodHook(typing_extensions.Protocol):
+    def __call__(self, __ctx: mypy.plugin.MethodContext) -> mypy.types.Type:
+        ...
+
+
+class AttributeHook(typing_extensions.Protocol):
+    def __call__(self, __ctx: mypy.plugin.AttributeContext) -> mypy.types.Type:
+        ...
+
+
+class ClassDefHook(typing_extensions.Protocol):
+    def __call__(self, __ctx: mypy.plugin.ClassDefContext) -> None:
+        ...

asyncpg/mypy/hooks.py

@@ -0,0 +1,108 @@
+import mypy.nodes
+import mypy.plugin
+import mypy.types
+import typing
+
+from . import common
+from . import utils
+
+
+def mark_record(ctx: mypy.plugin.ClassDefContext) -> None:
+    if ctx.cls.info.fullname == common.RECORD_NAME:
+        return
+
+    if ctx.cls.info.has_base(common.RECORD_NAME):
+        utils.mark_record(ctx.cls.info)
+
+
+def record_final_attributes(ctx: mypy.plugin.ClassDefContext) -> None:
+    if utils.is_record(ctx.cls.info):
+        for name, value in ctx.cls.info.names.items():
+            # set all properties as Final so they can't be set
+            if isinstance(value.node, mypy.nodes.Var):
+                value.node.is_final = True
+                value.node.final_set_in_init = True
+                value.node.final_unset_in_class = True
+
+
+def record_attribute(ctx: mypy.plugin.AttributeContext) \
+        -> mypy.types.Type:
+    if isinstance(ctx.type, mypy.types.Instance) and \
+            utils.is_record(ctx.type.type):
+        assert isinstance(ctx.context, mypy.nodes.MemberExpr)
+        symbol = ctx.type.type.get(ctx.context.name)
+
+        if symbol is None:
+            ctx.api.fail('"{}" has no attribute "{}"'
+                         .format(ctx.type.type.name, ctx.context.name),
+                         ctx.context)
+
+    return ctx.default_attr_type
+
+
+def record_getitem(ctx: mypy.plugin.MethodContext) \
+        -> mypy.types.Type:
+    if isinstance(ctx.type, mypy.types.Instance):
+        arg = ctx.arg_types[0][0]
+
+        if arg is not None and isinstance(arg, mypy.types.Instance) and \
+                isinstance(arg.last_known_value, mypy.types.LiteralType):
+            value = arg.last_known_value.value
+            names = utils.get_record_field_names(ctx.type.type.defn)
+            name = None  # type: typing.Optional[str]
+
+            if isinstance(value, int) and value < len(names):
+                name = names[value]
+            elif isinstance(value, str) and value in names:
+                name = value
+
+            if name is None:
+                ctx.api.fail('Unexpected key "{}" for record "{}'
+                             .format(value, ctx.type.type.name),
+                             ctx.context)
+            else:
+                node = ctx.type.type.get(name)
+
+                if node is not None and node.type is not None:
+                    return node.type
+
+    return ctx.default_return_type
+
+
+def record_get(ctx: mypy.plugin.MethodContext) \
+        -> mypy.types.Type:
+    if ctx.arg_names[0][0] is not None or \
+            len(ctx.arg_names) > 1 and \
+            ctx.arg_names[1][0] is not None:
+        ctx.api.fail('get() takes no keyword arguments', ctx.context)
+    elif isinstance(ctx.type, mypy.types.Instance):
+        arg = utils.get_argument_type_by_name(ctx, 'key')
+        default_arg = utils.get_argument_type_by_name(ctx, 'default')
+
+        if arg and isinstance(arg, mypy.types.Instance) and \
+                isinstance(arg.last_known_value, mypy.types.LiteralType):
+            value = arg.last_known_value.value
+            names = utils.get_record_field_names(ctx.type.type.defn)
+            name = None  # type: typing.Optional[str]
+
+            if isinstance(value, str) and value in names:
+                name = value
+
+            if name is None:
+                ctx.api.fail('Unexpected key "{}" for record "{}'
+                             .format(value, ctx.type.type.name),
+                             ctx.context)
+            else:
+                node = ctx.type.type.get(name)
+
+                assert node is not None
+
+                if node.type is not None:
+                    if default_arg is not None:
+                        return mypy.types.UnionType([node.type, default_arg],
+                                                    ctx.context.line,
+                                                    ctx.context.column)
+                    else:
+                        return node.type
+
+    return ctx.default_return_type