feat: Use model._fields directly.

Author: abawchen

graphene_mongo/tests/test_utils.py

@@ -5,14 +5,18 @@
 from ..utils import (
     get_model_fields, is_valid_mongoengine_model
 )
-from .models import Reporter
+from .models import Reporter, Child
 
 def test_get_model_fields_no_duplication():
     reporter_fields = get_model_fields(Reporter)
     reporter_name_set = set(reporter_fields)
     assert len(reporter_fields) == len(reporter_name_set)
 
 
+def test_get_base_model_fields():
+    child_fields = get_model_fields(Child)
+    assert all(field in set(child_fields) for field in ['bar', 'baz'])
+
 def test_is_valid_mongoengine_mode():
     assert is_valid_mongoengine_model(Reporter)
 

graphene_mongo/utils.py

@@ -2,20 +2,16 @@
 import mongoengine
 
 from collections import OrderedDict
-from mongoengine.base.fields import BaseField
 
 
 def get_model_fields(model, excluding=None):
     if excluding is None:
         excluding = []
     attributes = dict()
-    for attr_name in vars(model):
+    for attr_name, attr in model._fields.items():
         if attr_name in excluding:
             continue
-        attr = getattr(model, attr_name)
-        if isinstance(attr, BaseField):
-            attributes[attr_name] = attr
-
+        attributes[attr_name] = attr
     return OrderedDict(sorted(attributes.items()))