3
\$\begingroup\$

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.

\$\endgroup\$
5
\$\begingroup\$

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)
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.