switch_collection alias context manager. Example :: class Group(Document): name = StringField() Group(name='test').save() # Saves in the default db with switch_collection(Group, 'group1') as Group: Group(name='hello testdb!').save() # Saves i
| 88 | |
| 89 | |
| 90 | class switch_collection: |
| 91 | """switch_collection alias context manager. |
| 92 | |
| 93 | Example :: |
| 94 | |
| 95 | class Group(Document): |
| 96 | name = StringField() |
| 97 | |
| 98 | Group(name='test').save() # Saves in the default db |
| 99 | |
| 100 | with switch_collection(Group, 'group1') as Group: |
| 101 | Group(name='hello testdb!').save() # Saves in group1 collection |
| 102 | """ |
| 103 | |
| 104 | def __init__(self, cls, collection_name): |
| 105 | """Construct the switch_collection context manager. |
| 106 | |
| 107 | :param cls: the class to change the registered db |
| 108 | :param collection_name: the name of the collection to use |
| 109 | """ |
| 110 | self.cls = cls |
| 111 | self.ori_collection = cls._get_collection() |
| 112 | self.ori_get_collection_name = cls._get_collection_name |
| 113 | self.collection_name = collection_name |
| 114 | |
| 115 | def __enter__(self): |
| 116 | """Change the _get_collection_name and clear the cached collection.""" |
| 117 | |
| 118 | @classmethod |
| 119 | def _get_collection_name(cls): |
| 120 | return self.collection_name |
| 121 | |
| 122 | self.cls._get_collection_name = _get_collection_name |
| 123 | self.cls._collection = None |
| 124 | return self.cls |
| 125 | |
| 126 | def __exit__(self, t, value, traceback): |
| 127 | """Reset the collection.""" |
| 128 | self.cls._collection = self.ori_collection |
| 129 | self.cls._get_collection_name = self.ori_get_collection_name |
| 130 | |
| 131 | |
| 132 | @contextlib.contextmanager |
no outgoing calls