From d0e0b291df4d8c1501f5caad491b758a8f87ef64 Mon Sep 17 00:00:00 2001
From: Harry Marr <harry.marr@gmail.com>
Date: Thu, 25 Feb 2010 17:20:52 +0000
Subject: [PATCH] Implementation and tests for exec_js field substitution

---
 mongoengine/queryset.py | 18 ++++++++++++++
 tests/queryset.py       | 52 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+)

diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py
index 563c25ac..3ac8a37a 100644
--- a/mongoengine/queryset.py
+++ b/mongoengine/queryset.py
@@ -1,6 +1,7 @@
 from connection import _get_db
 
 import pymongo
+import re
 import copy
 
 
@@ -593,6 +594,21 @@ class QuerySet(object):
     def __iter__(self):
         return self
 
+    def _sub_js_fields(self, code):
+        """When fields are specified with [~fieldname] syntax, where 
+        *fieldname* is the Python name of a field, *fieldname* will be 
+        substituted for the MongoDB name of the field (specified using the
+        :attr:`name` keyword argument in a field's constructor).
+        """
+        def field_sub(match):
+            # Extract just the field name, and look up the field objects
+            field_name = match.group(1).split('.')
+            fields = QuerySet._lookup_field(self._document, field_name)
+            # Substitute the correct name for the field into the javascript
+            return '["%s"]' % fields[-1].name
+
+        return re.sub('\[\s*~([A-z_][A-z_0-9.]+?)\s*\]', field_sub, code)
+
     def exec_js(self, code, *fields, **options):
         """Execute a Javascript function on the server. A list of fields may be
         provided, which will be translated to their correct names and supplied
@@ -608,6 +624,8 @@ class QuerySet(object):
         :param options: options that you want available to the function 
             (accessed in Javascript through the ``options`` object)
         """
+        code = self._sub_js_fields(code)
+
         fields = [QuerySet._translate_field_name(self._document, f)
                   for f in fields]
         collection = self._document._meta['collection']
diff --git a/tests/queryset.py b/tests/queryset.py
index 61cb0537..d287efad 100644
--- a/tests/queryset.py
+++ b/tests/queryset.py
@@ -387,6 +387,58 @@ class QuerySetTest(unittest.TestCase):
 
         BlogPost.drop_collection()
 
+    def test_exec_js_field_sub(self):
+        """Ensure that field substitutions occur properly in exec_js functions.
+        """
+        class Comment(EmbeddedDocument):
+            content = StringField(name='body')
+
+        class BlogPost(Document):
+            name = StringField(name='doc-name')
+            comments = ListField(EmbeddedDocumentField(Comment), name='cmnts')
+
+        BlogPost.drop_collection()
+
+        comments1 = [Comment(content='cool'), Comment(content='yay')]
+        post1 = BlogPost(name='post1', comments=comments1)
+        post1.save()
+
+        comments2 = [Comment(content='nice stuff')]
+        post2 = BlogPost(name='post2', comments=comments2)
+        post2.save()
+
+        code = """
+        function getComments() {
+            var comments = [];
+            db[collection].find(query).forEach(function(doc) {
+                var docComments = doc[~comments];
+                for (var i = 0; i < docComments.length; i++) {
+                    comments.push({
+                        'document': doc[~name],
+                        'comment': doc[~comments][i][~comments.content]
+                    });
+                }
+            });
+            return comments;
+        }
+        """
+        
+        sub_code = BlogPost.objects._sub_js_fields(code)
+        code_chunks = ['doc["cmnts"];', 'doc["doc-name"],', 
+                       'doc["cmnts"][i]["body"]']
+        for chunk in code_chunks:
+            self.assertTrue(chunk in sub_code)
+
+        results = BlogPost.objects.exec_js(code)
+        expected_results = [
+            {u'comment': u'cool', u'document': u'post1'}, 
+            {u'comment': u'yay', u'document': u'post1'}, 
+            {u'comment': u'nice stuff', u'document': u'post2'},
+        ]
+        self.assertEqual(results, expected_results)
+
+        BlogPost.drop_collection()
+
     def test_delete(self):
         """Ensure that documents are properly deleted from the database.
         """