From 458930c74c8e2382a8bc931540254a5c38e84cba Mon Sep 17 00:00:00 2001
From: Anthony <ant@getme.co.uk>
Date: Thu, 16 Jan 2020 16:05:07 +0000
Subject: [PATCH] Added support for with_options method for Frames.

---
 .gitignore            |  1 +
 mongoframes/frames.py | 29 ++++++++++++++++++++++++++---
 pip-selfcheck.json    |  1 -
 setup.py              |  4 ++--
 tests/test_frames.py  | 23 ++++++++++++++++++++++-
 5 files changed, 51 insertions(+), 7 deletions(-)
 delete mode 100644 pip-selfcheck.json

diff --git a/.gitignore b/.gitignore
index edeeec1..8ed29d1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -19,3 +19,4 @@ pyvenv.cfg
 .tox
 .cache
 .pytest_cache/
+.pip-selfcheck.json
diff --git a/mongoframes/frames.py b/mongoframes/frames.py
index 3da9be9..b10cd57 100644
--- a/mongoframes/frames.py
+++ b/mongoframes/frames.py
@@ -1,3 +1,5 @@
+from contextlib import contextmanager
+
 from blinker import signal
 from bson.objectid import ObjectId
 from copy import deepcopy
@@ -430,12 +432,12 @@ def count(cls, filter=None, **kwargs):
 
         if filter:
             return cls.get_collection().count_documents(
-                to_refs(filter), 
+                to_refs(filter),
                 **kwargs
             )
         else:
             return cls.get_collection().estimated_document_count(**kwargs)
-            
+
     @classmethod
     def ids(cls, filter=None, **kwargs):
         """Return a list of Ids for documents matching the filter"""
@@ -779,7 +781,11 @@ def stop_listening(cls, event, func):
     @classmethod
     def get_collection(cls):
         """Return a reference to the database collection for the class"""
-        return getattr(cls.get_db(), cls._collection)
+        return getattr(
+            cls,
+            '_collection_context',
+            getattr(cls.get_db(), cls._collection)
+        )
 
     @classmethod
     def get_db(cls):
@@ -788,6 +794,23 @@ def get_db(cls):
             return getattr(cls._client, cls._db)
         return cls._client.get_default_database()
 
+    @classmethod
+    @contextmanager
+    def with_options(cls, **options):
+        existing_context = getattr(cls, '_collection_context', None)
+
+        try:
+            collection = getattr(cls.get_db(), cls._collection)
+            cls._collection_context = collection.with_options(**options)
+            yield cls._collection_context
+
+        finally:
+            if existing_context is None:
+                del cls._collection_context
+
+            else:
+                cls._collection_context = existing_context
+
 
 class SubFrame(_BaseFrame):
     """
diff --git a/pip-selfcheck.json b/pip-selfcheck.json
deleted file mode 100644
index 51fb13c..0000000
--- a/pip-selfcheck.json
+++ /dev/null
@@ -1 +0,0 @@
-{"last_check":"2018-04-12T15:08:33Z","pypi_version":"9.0.3"}
\ No newline at end of file
diff --git a/setup.py b/setup.py
index d73284c..92101ac 100644
--- a/setup.py
+++ b/setup.py
@@ -21,7 +21,7 @@
     # Versions should comply with PEP440.  For a discussion on single-sourcing
     # the version across setup.py and the project code, see
     # https://packaging.python.org/en/latest/single_source_version.html
-    version='1.3.4',
+    version='1.3.5',
     description='A fast unobtrusive MongoDB ODM for Python',
     long_description=long_description,
     long_description_content_type='text/markdown',
@@ -84,7 +84,7 @@
     install_requires=[
         'blinker>=1.4',
         'Faker>=0.7.18',
-        'pymongo>=3.8'
+        'pymongo>=3.9.0'
         ],
 
     # List additional groups of dependencies here (e.g. development
diff --git a/tests/test_frames.py b/tests/test_frames.py
index a885fa8..91e43c5 100644
--- a/tests/test_frames.py
+++ b/tests/test_frames.py
@@ -3,6 +3,7 @@
 from time import sleep
 from unittest.mock import Mock
 
+from pymongo import ReadPreference
 from mongoframes import *
 
 from tests.fixtures import *
@@ -749,4 +750,24 @@ def test_flattern_projection():
         'name': True,
         'inventory.gold': True,
         'inventory.secret_draw.gold': True
-    }
\ No newline at end of file
+    }
+
+def test_with_options():
+    """Flattern projection"""
+
+    collection = Dragon.get_collection()
+
+    with Dragon.with_options(read_preference=ReadPreference.SECONDARY):
+        assert Dragon.get_collection().read_preference \
+                == ReadPreference.SECONDARY
+
+        with Dragon.with_options(
+                read_preference=ReadPreference.PRIMARY_PREFERRED):
+
+            assert Dragon.get_collection().read_preference \
+                    == ReadPreference.PRIMARY_PREFERRED
+
+        assert Dragon.get_collection().read_preference \
+                == ReadPreference.SECONDARY
+
+    assert collection.read_preference == ReadPreference.PRIMARY