Make set_local_events handle BRANCH event

Author: markshannon

Lib/test/test_monitoring.py

@@ -1853,6 +1853,10 @@ def test_local(self):
         code = f1.__code__
         sys.monitoring.set_local_events(TEST_TOOL, code, E.PY_START)
         self.assertEqual(sys.monitoring.get_local_events(TEST_TOOL, code), E.PY_START)
+        sys.monitoring.set_local_events(TEST_TOOL, code, 0)
+        sys.monitoring.set_local_events(TEST_TOOL, code, E.BRANCH)
+        self.assertEqual(sys.monitoring.get_local_events(TEST_TOOL, code), E.BRANCH_NOT_TAKEN | E.BRANCH_TAKEN)
+        sys.monitoring.set_local_events(TEST_TOOL, code, 0)
         sys.monitoring.set_local_events(TEST_TOOL2, code, E.PY_START)
         self.assertEqual(sys.monitoring.get_local_events(TEST_TOOL2, code), E.PY_START)
         sys.monitoring.set_local_events(TEST_TOOL, code, 0)

Python/instrumentation.c

@@ -2264,6 +2264,10 @@ monitoring_set_local_events_impl(PyObject *module, int tool_id,
         return NULL;
     }
     event_set &= ~C_RETURN_EVENTS;
+    if (event_set & (1 << PY_MONITORING_EVENT_BRANCH)) {
+        event_set &= ~(1 << PY_MONITORING_EVENT_BRANCH);
+        event_set |= (1 << PY_MONITORING_EVENT_BRANCH_TAKEN) | (1 << PY_MONITORING_EVENT_BRANCH_NOT_TAKEN);
+    }
     if (event_set < 0 || event_set >= (1 << _PY_MONITORING_LOCAL_EVENTS)) {
         PyErr_Format(PyExc_ValueError, "invalid local event set 0x%x", event_set);
         return NULL;