Cupy context manager

cupy_context_manager.py

class CupyContextManager:
    def __init__(self):
        self.previous_np = None

    def __enter__(self):
        self.previous_np = globals().get("np")
        import cupy as np
        globals()["np"] = np
        return np

    def __exit__(self, exc_type, exc_value, exc_tb):
        if self.previous_np is not None:
            globals()["np"] = self.previous_np
        else:
            globals().pop('np', None)


with CupyContextManager() as np:
    print(np)