Support Falsey primary_keys (#1354)

This commit is contained in:
Ollie Ford 2016-12-03 21:10:05 +00:00 committed by Stefan Wójcik
parent 02fb3b9315
commit d8d98b6143
2 changed files with 18 additions and 3 deletions

View File

@ -472,7 +472,7 @@ class Document(BaseDocument):
Raises :class:`OperationError` if called on an object that has not yet Raises :class:`OperationError` if called on an object that has not yet
been saved. been saved.
""" """
if not self.pk: if self.pk is None:
if kwargs.get('upsert', False): if kwargs.get('upsert', False):
query = self.to_mongo() query = self.to_mongo()
if "_cls" in query: if "_cls" in query:
@ -604,7 +604,7 @@ class Document(BaseDocument):
elif "max_depth" in kwargs: elif "max_depth" in kwargs:
max_depth = kwargs["max_depth"] max_depth = kwargs["max_depth"]
if not self.pk: if self.pk is None:
raise self.DoesNotExist("Document does not exist") raise self.DoesNotExist("Document does not exist")
obj = self._qs.read_preference(ReadPreference.PRIMARY).filter( obj = self._qs.read_preference(ReadPreference.PRIMARY).filter(
**self._object_key).only(*fields).limit( **self._object_key).only(*fields).limit(
@ -655,7 +655,7 @@ class Document(BaseDocument):
def to_dbref(self): def to_dbref(self):
"""Returns an instance of :class:`~bson.dbref.DBRef` useful in """Returns an instance of :class:`~bson.dbref.DBRef` useful in
`__raw__` queries.""" `__raw__` queries."""
if not self.pk: if self.pk is None:
msg = "Only saved documents can have a valid dbref" msg = "Only saved documents can have a valid dbref"
raise OperationError(msg) raise OperationError(msg)
return DBRef(self.__class__._get_collection_name(), self.pk) return DBRef(self.__class__._get_collection_name(), self.pk)

View File

@ -3202,5 +3202,20 @@ class InstanceTest(unittest.TestCase):
self.assertEqual(b._instance, a) self.assertEqual(b._instance, a)
self.assertEqual(idx, 2) self.assertEqual(idx, 2)
def test_falsey_pk(self):
"""Ensure that we can create and update a document with Falsey PK.
"""
class Person(Document):
age = IntField(primary_key=True)
height = FloatField()
person = Person()
person.age = 0
person.height = 1.89
person.save()
person.update(set__height=2.0)
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()