Commits

James Crasta  committed 06a6d88

appengine.ndb converts KeyProperty properly. Fixes #159.

  • Participants
  • Parent commits 7be18f7

Comments (0)

Files changed (3)

File tests/ext_appengine/test_ndb.py

 
 
 class Book(ndb.Model):
-    author = ndb.KeyProperty(Author)
+    author = ndb.KeyProperty(kind=Author)
 
 
 class TestKeyPropertyField(TestCase):
 class TestModelForm(TestCase):
     EXPECTED_AUTHOR = [('name', TextField), ('city', TextField), ('age', IntegerField), ('is_admin', BooleanField)]
 
-    def test(self):
+    def test_author(self):
         form = model_form(Author)
         for (expected_name, expected_type), field in zip(self.EXPECTED_AUTHOR, form()):
             self.assertEqual(field.name, expected_name)
             self.assertEqual(type(field), expected_type)
+
+    def test_book(self):
+        authors = set(text_type(x.key.id()) for x in fill_authors(Author))
+        authors.add('__None')
+        form = model_form(Book)
+        keys = set()
+        for key, b, c in form().author.iter_choices():
+            keys.add(key)
+
+        self.assertEqual(authors, keys)

File tests/ext_appengine/tests.py

         self.assertEqual(form.is_admin.label.text, 'Administrative rights')
 
     def test_reference_property(self):
-        keys = ['__None']
+        keys = set(['__None'])
         for name in ['foo', 'bar', 'baz']:
             author = Author(name=name, age=26)
             author.put()
-            keys.append(str(author.key()))
+            keys.add(str(author.key()))
 
         form_class = model_form(Book)
         form = form_class()
 
-        i = 0
         for key, name, value in form.author.iter_choices():
-            self.assertEqual(key, keys[i])
-            i += 1
+            assert key in keys
+            keys.remove(key)
+
+        assert not keys
 
 
 class TestFields(TestCase):

File wtforms/ext/appengine/ndb.py

 
 """
 from wtforms import Form, validators, fields as f
+from wtforms.compat import string_types
 from wtforms.ext.appengine.fields import GeoPtPropertyField, KeyPropertyField, StringListPropertyField, IntegerListPropertyField
 
 
 
     def convert_KeyProperty(self, model, prop, kwargs):
         """Returns a form field for a ``ndb.KeyProperty``."""
-        kwargs['reference_class'] = prop._reference_class
+        if 'reference_class' not in kwargs:
+            try:
+                reference_class = prop._kind
+            except AttributeError:
+                reference_class = prop._reference_class
+
+            if isinstance(reference_class, string_types):
+                # reference class is a string, try to retrieve the model object.
+                mod = __import__(model.__module__, None, None, [reference_class], 0)
+                reference_class = getattr(mod, reference_class)
+            kwargs['reference_class'] = reference_class
         kwargs.setdefault('allow_blank', not prop._required)
         return KeyPropertyField(**kwargs)