from rainbowadn.instrument import *


class Print(Instrumentation):
    def __init__(self, target, methodname: str, msg: str):
        super().__init__(target, methodname)
        self.msg = msg

    def instrument(self, method, *args, **kwargs):
        print(self.msg, end=" ")
        return method(*args, **kwargs)


class C:
    @classmethod
    def m(cls):
        print("m")


with Print(Instrumentation, "__exit__", "exit"):
    print1 = Print(C, "m", "1")
    print2 = Print(C, "m", "2")
    print3 = Print(C, "m", "3")
    C.m()
    print1.__enter__()
    C.m()
    print2.__enter__()
    C.m()
    print3.__enter__()
    C.m()
    print1.__exit__(None, None, None)
    C.m()
    print2.__exit__(None, None, None)
    C.m()
    print3.__exit__(None, None, None)
    C.m()