import pymongo


class Client:

    def __init__(self, uri, dbname, collection=None):
        self.client = pymongo.MongoClient(uri)
        self.db = self.client[dbname]
        self.collection = collection


    def __enter__(self):
        if self.collection:
            obj = getattr(self.db, self.collection)
        else:
            obj = self.db

        return obj


    def __exit__(self, *args, **kwargs):
        self.client.close()