add plugin support for ArrayField base_field

mypy_django_plugin/plugins/postgres_fields.py

@@ -1,14 +1,24 @@
-from mypy.plugin import Plugin, ClassDefContext
+from typing import Optional, Callable
+
+from mypy.plugin import Plugin, FunctionContext
+from mypy.types import Type
 
 
-def determine_type_of_array_field(context: ClassDefContext) -> None:
-    pass
+def determine_type_of_array_field(ctx: FunctionContext) -> Type:
+    assert 'base_field' in ctx.context.arg_names
+    base_field_arg_index = ctx.context.arg_names.index('base_field')
+    base_field_arg_type = ctx.arg_types[base_field_arg_index][0]
+
+    return ctx.api.named_generic_type(ctx.context.callee.fullname,
+                                      args=[base_field_arg_type.type.names['__get__'].type.ret_type])
 
 
 class PostgresFieldsPlugin(Plugin):
-    def get_base_class_hook(self, fullname: str
-                            ):
-        return determine_type_of_array_field
+    def get_function_hook(self, fullname: str
+                          ) -> Optional[Callable[[FunctionContext], Type]]:
+        if fullname == 'django.contrib.postgres.fields.array.ArrayField':
+            return determine_type_of_array_field
+        return None
 
 
 def plugin(version):