Skip to content

Commit

Permalink
Use correct class when deserializing a PolyModel entity.
Browse files Browse the repository at this point in the history
  • Loading branch information
Chris Rossi committed Sep 9, 2019
1 parent 7bdd01f commit b4639c3
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 5 deletions.
8 changes: 7 additions & 1 deletion google/cloud/ndb/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -527,7 +527,13 @@ def _entity_from_ds_entity(ds_entity, model_class=None):
Returns:
.Model: The deserialized entity.
"""
model_class = model_class or Model._lookup_model(ds_entity.kind)
class_key = ds_entity.get("class")
if class_key:
kind = class_key[-1]
else:
kind = ds_entity.kind

model_class = model_class or Model._lookup_model(kind)
entity = model_class()

# Check if we are dealing with a PolyModel, and if so get correct subclass.
Expand Down
11 changes: 7 additions & 4 deletions tests/system/test_crud.py
Original file line number Diff line number Diff line change
Expand Up @@ -853,21 +853,24 @@ class SomeKind(ndb.Expando):
@pytest.mark.usefixtures("client_context")
def test_insert_polymodel(dispose_of):
class Animal(ndb.PolyModel):
pass
one = ndb.StringProperty()

class Feline(Animal):
pass
two = ndb.StringProperty()

class Cat(Feline):
pass
three = ndb.StringProperty()

entity = Cat()
entity = Cat(one="hello", two="dad", three="i'm in jail")
key = entity.put()

retrieved = key.get()

assert isinstance(retrieved, Animal)
assert isinstance(retrieved, Cat)
assert retrieved.one == "hello"
assert retrieved.two == "dad"
assert retrieved.three == "i'm in jail"

dispose_of(key._key)

Expand Down

0 comments on commit b4639c3

Please sign in to comment.