from rainbowadn.testing.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()