import pymongo
from .client import TracedMongoClient
# Original Client class
_MongoClient = pymongo.MongoClient
def patch():
setattr(pymongo, 'MongoClient', TracedMongoClient)
def unpatch():
setattr(pymongo, 'MongoClient', _MongoClient)