Fix an issue where encoding a simple ndb object would write out a key…

… instead.
hydralabs · Feb 6, 2015 · c2b6bb1 · c2b6bb1
1 parent 97e1336
commit c2b6bb1
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 11 deletions.
diff --git a/pyamf/adapters/_google_appengine_ext_ndb.py b/pyamf/adapters/_google_appengine_ext_ndb.py
@@ -235,10 +235,16 @@ def getDecodableAttributes(self, obj, attrs, codec=None):
             key = ndb.Key(urlsafe=key)
 
         if self.model_properties:
-            property_attrs = [k for k in attrs if k in self.model_properties]
+            property_attrs = [
+                k for k in attrs if k in self.decodable_properties
+            ]
 
             for name in property_attrs:
-                prop = self.model_properties[name]
+                prop = self.model_properties.get(name, None)
+
+                if not prop:
+                    continue
+
                 value = attrs[name]
 
                 if not prop._repeated:
@@ -278,7 +284,12 @@ def getEncodableAttributes(self, obj, codec=None):
                 del attrs[k]
 
         if self.model_properties:
-            for name, prop in self.model_properties.iteritems():
+            for name in self.encodable_properties:
+                prop = self.model_properties.get(name, None)
+
+                if not prop:
+                    continue
+
                 attrs[name] = self.getAttribute(obj, name, codec=codec)
 
                 if prop._repeated:
@@ -297,7 +308,7 @@ def getEncodableAttributes(self, obj, codec=None):
                     attrs[name],
                 )
 
-        attrs[self.KEY_ATTR] = obj.key
+        attrs[self.KEY_ATTR] = unicode(obj.key.urlsafe()) if obj.key else None
 
         return attrs
 
@@ -356,13 +367,20 @@ def encode_ndb_key(key, encoder=None):
     When encountering an L{ndb.Key} instance, find the entity in the datastore
     and encode that.
     """
-    if not key:
-        return key
+    klass = ndb.Model._kind_map.get(key.kind())
 
-    if key.id():
-        return key.urlsafe()
+    referenced_object = _get_by_class_key(
+        encoder,
+        klass,
+        key,
+    )
 
-    return None
+    if not referenced_object:
+        encoder.writeNull(None)
+
+        return
+
+    encoder.writeObject(referenced_object)
 
 
 def _get_by_class_key(codec, klass, key, obj=None):
@@ -416,6 +434,14 @@ def decode_int_property(prop, value):
     return value
 
 
+@adapter_models.register_property_encoder(ndb.KeyProperty)
+def encode_key_property(prop, value):
+    if not hasattr(value, 'urlsafe'):
+        return value
+
+    return value.urlsafe()
+
+
 @adapter_models.register_property_encoder(ndb.TimeProperty)
 def encode_time_property(prop, value):
     # PyAMF supports datetime.datetime objects and won't decide what date to

diff --git a/pyamf/adapters/tests/google/test_ndb.py b/pyamf/adapters/tests/google/test_ndb.py
@@ -124,11 +124,11 @@ def test_key(self):
         key = ndb.Key('SimpleEntity', 'bar')
 
         self.assertEncodes(key, (
-            b'\x02\x002agx0ZXN0YmVkLXRlc3RyFQsSDFNpbXBsZUVudGl0eSIDYmFyDA'
+            b'\x05',
         ), encoding=pyamf.AMF0)
 
         self.assertEncodes(key, (
-            b'\x06eagx0ZXN0YmVkLXRlc3RyFQsSDFNpbXBsZUVudGl0eSIDYmFyDA'
+            b'\x01'
         ), encoding=pyamf.AMF3)
 
     def test_query(self):