nested class Meta support

mypy_django_plugin/plugins/fields.py

@@ -0,0 +1,41 @@
+from typing import Iterator, List, cast
+
+from mypy.nodes import ClassDef, AssignmentStmt, CallExpr
+from mypy.plugin import FunctionContext, ClassDefContext
+from mypy.semanal import SemanticAnalyzerPass2
+from mypy.types import Type, Instance
+
+from mypy_django_plugin.plugins.related_fields import add_new_var_node_to_class
+
+
+def determine_type_of_array_field(ctx: FunctionContext) -> Type:
+    if 'base_field' not in ctx.arg_names:
+        return ctx.default_return_type
+
+    base_field_arg_type = ctx.arg_types[ctx.arg_names.index('base_field')][0]
+    return ctx.api.named_generic_type(ctx.context.callee.fullname,
+                                      args=[base_field_arg_type.type.names['__get__'].type.ret_type])
+
+
+def get_assignments(klass: ClassDef) -> List[AssignmentStmt]:
+    stmts = []
+    for stmt in klass.defs.body:
+        if not isinstance(stmt, AssignmentStmt):
+            continue
+        if len(stmt.lvalues) > 1:
+            # not supported yet
+            continue
+        stmts.append(stmt)
+    return stmts
+
+
+def add_int_id_attribute_if_primary_key_true_is_not_present(ctx: ClassDefContext) -> None:
+    api = cast(SemanticAnalyzerPass2, ctx.api)
+    for stmt in get_assignments(ctx.cls):
+        if (isinstance(stmt.rvalue, CallExpr)
+                and 'primary_key' in stmt.rvalue.arg_names
+                    and api.parse_bool(stmt.rvalue.args[stmt.rvalue.arg_names.index('primary_key')])):
+            break
+    else:
+        add_new_var_node_to_class(ctx.cls.info, 'id', api.builtin_type('builtins.int'))
+