RAII style API wrapper for PyMongo

Question

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.

Answer 1

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)