# RAII style API wrapper for PyMongo

I just watched Raymond Hettinger's talk on making Python more Pythonic and realized I should be putting a lot of his ideas into practice, particularly wrapping API's in a class that makes everything simpler and easy to use. Here's what I'm doing to wrap PyMongo:

from pymongo import MongoClient

class MongoDB(object):
"""Provides a RAII wrapper for PyMongo db connections.
Available collection functions limited to those in
attributes_to_pass. Number of simultaneous connection
users also tracked. """

attributes_to_pass = ["update", "insert", "count"]
client = None
num_users = 0

def __init__(self, db, collection):
MongoDB.client = MongoDB.client or MongoClient()
self.collection = MongoDB.client[db][collection]
MongoDB.num_users += 1

def __enter__(self, *args):
return self

def __exit__(self, type, value, traceback):
MongoDB.num_users -= 1
if MongoDB.num_users is 0:
self.client.close()

def __getattr__(self, attr):
if attr in MongoDB.attributes_to_pass:
return getattr(self.collection, attr)
else:
return getattr(self, attr)

def main():
with MongoDB(db = "db1", collection = "c1") as m:
print(m.count())
m.update({"jello":5} , {"hello":"you"}, upsert = True)
print(m.count())
m.insert({"joe":6})
with MongoDB(db ='db1', collection = 'c2') as j:
j.insert({"joe":6})
print(j.count())

if __name__ == "__main__":
main()


I would really appreciate all suggestions on how to make this better.

Great that you provided a docstring! You made one minor formatting mistake though, there should be a blank line between your brief single line summary and the rest of the docstring. It's recommended in the style guide partially for readability and partially for script parsers.

class MongoDB(object):
"""Provides a RAII wrapper for PyMongo db connections.

Available collection functions limited to those in
attributes_to_pass. Number of simultaneous connection
users also tracked. """


It seems like attributes_to_pass is actually a constant. If that is the case it should be UPPER_SNAKE_CASE to be clear that it is, especially when mixed with attributes that often change like client and num_users. It should also be a tuple, as tuples are immutable so will not change unless updated in the source code.

ATTRIBUTES_TO_PASS = ("update", "insert", "count")


Also, when you call on these constants you refer to the class name. While that's possible I'd personally say it's less readable. I had thought you were calling on a class you imported for a second. You can just pass self instead of the classname.

def __getattr__(self, attr):
if attr in self.ATTRIBUTES_TO_PASS:
return getattr(self.collection, attr)
else:
return getattr(self, attr)