From 84f9e44b6cb6a3d8ccc1534a4d5bdec2c8e578eb Mon Sep 17 00:00:00 2001
From: Ross Lawley <ross.lawley@gmail.com>
Date: Mon, 5 Dec 2011 04:16:57 -0800
Subject: [PATCH] Fixed GridFS documents can now be pickled

Refs #135 #381
---
 docs/changelog.rst    |  1 +
 mongoengine/fields.py | 26 ++++++++++++++++++++++----
 tests/fixtures.py     |  1 +
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/docs/changelog.rst b/docs/changelog.rst
index 6dd9dbc0..81b34306 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -5,6 +5,7 @@ Changelog
 Changes in dev
 ==============
 
+- Fixed GridFS documents can now be pickled
 - Added Now raises an InvalidDocumentError when declaring multiple fields with the same db_field
 - Added InvalidQueryError when calling with_id with a filter
 - Added support for DBRefs in distinct()
diff --git a/mongoengine/fields.py b/mongoengine/fields.py
index 68d7acee..914cb925 100644
--- a/mongoengine/fields.py
+++ b/mongoengine/fields.py
@@ -786,18 +786,25 @@ class GridFSProxy(object):
     .. versionchanged:: 0.6 - added collection name param
     """
 
+    _fs = None
+
     def __init__(self, grid_id=None, key=None,
                  instance=None,
                  db_alias=DEFAULT_CONNECTION_NAME,
                  collection_name='fs'):
-        self.fs = gridfs.GridFS(get_db(db_alias), collection_name)  # Filesystem instance
-        self.newfile = None                 # Used for partial writes
-        self.grid_id = grid_id              # Store GridFS id for file
-        self.gridout = None
+        self.grid_id = grid_id                  # Store GridFS id for file
         self.key = key
         self.instance = instance
+        self.db_alias = db_alias
+        self.collection_name = collection_name
+        self.newfile = None                     # Used for partial writes
+        self.gridout = None
 
     def __getattr__(self, name):
+        attrs = ('_fs', 'grid_id', 'key', 'instance', 'db_alias',
+                 'collection_name', 'newfile', 'gridout')
+        if name in attrs:
+            return self.__getattribute__(name)
         obj = self.get()
         if name in dir(obj):
             return getattr(obj, name)
@@ -809,6 +816,17 @@ class GridFSProxy(object):
     def __nonzero__(self):
         return bool(self.grid_id)
 
+    def __getstate__(self):
+        self_dict = self.__dict__
+        self_dict['_fs'] = None
+        return self_dict
+
+    @property
+    def fs(self):
+        if not self._fs:
+            self._fs = gridfs.GridFS(get_db(self.db_alias), self.collection_name)
+        return self._fs
+
     def get(self, id=None):
         if id:
             self.grid_id = id
diff --git a/tests/fixtures.py b/tests/fixtures.py
index 32081fe1..a97e8eba 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -12,6 +12,7 @@ class PickleTest(Document):
     string = StringField(choices=(('One', '1'), ('Two', '2')))
     embedded = EmbeddedDocumentField(PickleEmbedded)
     lists = ListField(StringField())
+    photo = FileField()
 
 
 class Mixin(object):