From 90e27cc87dcc02b1f9ef625fab6d9230d91a21ad Mon Sep 17 00:00:00 2001
From: Harry Marr <harry.marr@gmail.com>
Date: Thu, 19 Nov 2009 20:28:43 +0000
Subject: [PATCH] find_one now supports using just an object id

---
 mongomap/collection.py | 23 ++++++++++++++++++-----
 tests/collection.py    |  4 ++++
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/mongomap/collection.py b/mongomap/collection.py
index 3f0ae8f8..83cea585 100644
--- a/mongomap/collection.py
+++ b/mongomap/collection.py
@@ -1,5 +1,7 @@
 from connection import _get_db
 
+import pymongo
+
 
 class QuerySet(object):
     """A set of results returned from a query. Wraps a MongoDB cursor, 
@@ -74,15 +76,26 @@ class CollectionManager(object):
         return mongo_query
 
     def find(self, **query):
-        """Query the collection for document matching the provided query.
+        """Query the collection for documents matching the provided query.
         """
         query = self._transform_query(**query)
         query['_types'] = self._document._class_name
         return QuerySet(self._document, self._collection.find(query))
 
-    def find_one(self, **query):
+    def find_one(self, object_id=None, **query):
         """Query the collection for document matching the provided query.
         """
-        query = self._transform_query(**query)
-        query['_types'] = self._document._class_name
-        return self._document._from_son(self._collection.find_one(query))
+        if object_id:
+            # Use just object_id if provided
+            if not isinstance(object_id, pymongo.objectid.ObjectId):
+                object_id = pymongo.objectid.ObjectId(object_id)
+            query = object_id
+        else:
+            # Otherwise, use the query provided
+            query = self._transform_query(**query)
+            query['_types'] = self._document._class_name
+
+        result = self._collection.find_one(query)
+        if result is not None:
+            result = self._document._from_son(result)
+        return result
diff --git a/tests/collection.py b/tests/collection.py
index f5b69a3c..1b7f1026 100644
--- a/tests/collection.py
+++ b/tests/collection.py
@@ -96,6 +96,10 @@ class CollectionManagerTest(unittest.TestCase):
 
         person = self.Person.objects.find_one(age__lt=30)
         self.assertEqual(person.name, "User A")
+        
+        # Find a document using just the object id
+        person = self.Person.objects.find_one(person1._id)
+        self.assertEqual(person.name, "User A")
 
     def test_find_embedded(self):
         """Ensure that an embedded document is properly returned from a query.