How to use the yappi.set_context_id_callback function in yappi

def test_callback(self):
        self.context_id = 0
        self.context_name = 'a'
        yappi.set_context_id_callback(lambda: self.context_id)
        yappi.set_context_name_callback(lambda: self.context_name)
        yappi.start()
        a()
        self.context_id = 1
        self.context_name = 'b'
        a()

        # Re-schedule context 0.
        self.context_id = 0
        self.context_name = 'a'
        a()
        yappi.stop()

        threadstats = yappi.get_thread_stats().sort('name', 'ascending')
        self.assertEqual(2, len(threadstats))
        self.assertEqual(0, threadstats[0].id)

def test_callback_function_int_return_overflow(self):
        # this test is just here to check if any errors are generated, as the err
        # is printed in C side, I did not include it here. THere are ways to test
        # this deterministically, I did not bother
        import ctypes

        def _unsigned_overflow_margin():
            return 2**(ctypes.sizeof(ctypes.c_void_p) * 8) - 1

        def foo():
            pass

        #with utils.captured_output() as (out, err):
        yappi.set_context_id_callback(_unsigned_overflow_margin)
        yappi.set_tag_callback(_unsigned_overflow_margin)
        yappi.start()
        foo()

def test_callback_non_integer(self):
        self.callback_count = 0

        def callback():
            self.callback_count += 1
            return None  # Supposed to return an integer.

        yappi.set_context_id_callback(callback)
        yappi.start()
        a()
        a()
        yappi.stop()

        # Callback was cleared after first error.
        self.assertEqual(1, self.callback_count)

def test_clear_callback(self):
        self.callback_count = 0

        def callback():
            self.callback_count += 1
            return 1

        yappi.set_context_id_callback(callback)
        yappi.start()
        a()
        yappi.set_context_id_callback(None)
        old_callback_count = self.callback_count
        a()
        yappi.stop()

        self.assertEqual(old_callback_count, self.callback_count)

def test_pause_resume(self):
        yappi.set_context_id_callback(lambda: self.context_id)
        yappi.set_clock_type('wall')

        # Start in context 0.
        self.context_id = 0
        yappi.start()
        time.sleep(0.08)

        # Switch to context 1.
        self.context_id = 1
        time.sleep(0.05)

        # Switch back to context 0.
        self.context_id = 0
        time.sleep(0.07)

        yappi.stop()

def test_callback_error(self):
        self.callback_count = 0

        def callback():
            self.callback_count += 1
            raise Exception('callback error')

        yappi.set_context_id_callback(callback)
        yappi.start()
        a()
        a()
        yappi.stop()

        # Callback was cleared after first error.
        self.assertEqual(1, self.callback_count)

def _capture_profile(fname=''):
    if not fname:
        yappi.set_clock_type('cpu')
        # We need to set context to greenlet to profile greenlets
        # https://bitbucket.org/sumerc/yappi/pull-requests/3
        yappi.set_context_id_callback(
            lambda: id(greenlet.getcurrent()))
        yappi.set_context_name_callback(
            lambda: greenlet.getcurrent().__class__.__name__)
        yappi.start()
    else:
        yappi.stop()
        stats = yappi.get_func_stats()
        # User should provide filename. This file with a suffix .prof
        # will be created in temp directory.
        try:
            stats_file = os.path.join(tempfile.gettempdir(), fname + '.prof')
            stats.save(stats_file, "pstat")
        except Exception as e:
            print("Error while saving the trace stats ", str(e))
        finally:
            yappi.clear_stats()

def start(builtins=False, profile_threads=True):
    """Start profiler."""
    # TODO: what about builtins False or profile_threads False?
    yappi.set_context_id_callback(lambda: id(greenlet.getcurrent()))
    yappi.start(builtins, profile_threads)