PolyModel with endpoints-proto-datastore

from google.appengine.ext import ndb
from google.appengine.ext.ndb import polymodel
from protorpc import remote

import endpoints

from endpoints_proto_datastore.ndb import EndpointsModel
from endpoints_proto_datastore.ndb.model import _EndpointsQueryInfo


def _DowncastMessage(message, final_message_class):
  message_class = message.__class__

  downcasted_message = final_message_class()
  for field in final_message_class.all_fields():
    # KeyError allowed to happen, if the field is missing,
    # a downcast should be performed
    field_on_message = message_class.field_by_name(field.name)
    # Make sure fields are the same type
    if field_on_message.__class__ != field.__class__:
      raise TypeError('Downcasted field of the wrong type: %r. Should be %r.' %
                      (field_on_message.__class__, field.__class__))
    value = getattr(message, field.name)
    setattr(downcasted_message, field.name, value)
  return downcasted_message


class _PolyModelQueryInfo(_EndpointsQueryInfo):

  def _PopulateFilters(self):
    entity = self._entity
    for prop in entity._properties.itervalues():
      if isinstance(prop, polymodel._ClassKeyProperty):
        continue
      attr_name = prop._code_name
      current_value = getattr(entity, attr_name)

      # Only filter for non-null values
      if current_value is not None:
        self._AddFilter(prop == current_value)


class MyModel(EndpointsModel, polymodel.PolyModel):
  _message_fields_schema = ('attr1', 'created')

  attr1 = ndb.StringProperty()
  created = ndb.DateTimeProperty(auto_now_add=True)

  def __init__(self, *args, **kwargs):
    # Don't need to call both constructors since PolyModel doesn't define one
    # and descends from model.Model, a superclass of EndpointsModel
    super(MyModel, self).__init__(*args, **kwargs)
    self._endpoints_query_info = _PolyModelQueryInfo(self)

  @classmethod
  def ToMessageCollection(cls, items, collection_fields=None,
                          next_cursor=None):
    proto_model = cls.ProtoCollection(collection_fields=collection_fields)

    items_as_message = [item.ToMessage(fields=collection_fields)
                        for item in items]
    final_proto_class = cls.ProtoModel(fields=collection_fields)
    items_as_message = [_DowncastMessage(item, final_proto_class)
                        for item in items_as_message]

    result = proto_model(items=items_as_message)

    if next_cursor is not None:
      result.nextPageToken = next_cursor.to_websafe_string()

    return result


class NextModel(MyModel):
  _message_fields_schema = ('attr1', 'attr2', 'created')
  attr2 = ndb.StringProperty()


@endpoints.api(name='myapi', version='v1', description='My Little API')
class MyApi(remote.Service):

  @MyModel.method(path='mymodel', http_method='POST', name='mymodel.insert')
  def MyModelInsert(self, my_model):
    my_model.put()
    return my_model

  @MyModel.query_method(path='mymodels', name='mymodel.list')
  def MyModelList(self, query):
    return query

  @NextModel.method(path='nextmodel', name='nextmodel.insert')
  def NextModelInsert(self, my_model):
    my_model.put()
    return my_model

  @NextModel.query_method(path='nextmodels', name='nextmodel.list')
  def NextModelList(self, query):
    return query


application = endpoints.api_server([MyApi], restricted=False)