1 #include <torch/csrc/dynamo/framelocals_mapping.h> 2 3 #if IS_PYTHON_3_12_PLUS 4 #include <torch/csrc/dynamo/cpython_defs.h> 5 #include <torch/csrc/dynamo/cpython_includes.h> 6 #include <torch/csrc/dynamo/debug_macros.h> 7 #include <torch/csrc/utils/pybind.h> 8 9 #include <internal/pycore_code.h> 10 11 // Our own version of PyFrame_GetLocals. 12 // Also combines functionality from frame_init_get_vars and frame_get_var. 13 // PyFrame_GetLocals: 14 // https://github.com/python/cpython/blob/0325a8a8cdba6c091bcbbb3c995f3bf1d1217012/Objects/frameobject.c#L1213 15 // frame_init_get_vars: 16 // https://github.com/python/cpython/blob/0325a8a8cdba6c091bcbbb3c995f3bf1d1217012/Objects/frameobject.c#L1136 17 // frame_get_var: 18 // https://github.com/python/cpython/blob/0325a8a8cdba6c091bcbbb3c995f3bf1d1217012/Objects/frameobject.c#L1162 19 // PyFrame_GetLocals returns the frame locals dict. 20 // frame_init_get_vars initializes free variables from the closure. 21 // frame_get_var fetches the variable value from the frame given the index 22 // NOTE: hidden variables are not included. 23 // Returns a new reference. get_framelocals_mapping(_PyInterpreterFrame * frame)24PyObject* get_framelocals_mapping(_PyInterpreterFrame* frame) { 25 if (!frame->stacktop) { 26 return py::dict().release().ptr(); 27 } 28 29 PyCodeObject* co = F_CODE(frame); 30 py::dict mapping; 31 32 auto update_mapping = [&](int i, PyObject* value) { 33 _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i); 34 35 if (kind & CO_FAST_FREE && !(co->co_flags & CO_OPTIMIZED)) { 36 return; 37 } 38 if (kind & CO_FAST_HIDDEN) { 39 return; 40 } 41 42 if (kind & CO_FAST_FREE) { 43 CHECK(value != nullptr && PyCell_Check(value)); 44 value = PyCell_GET(value); 45 } 46 47 if (value != nullptr) { 48 py::str name = 49 py::cast<py::str>(PyTuple_GET_ITEM(co->co_localsplusnames, i)); 50 mapping[name] = py::cast<py::object>(value); 51 } 52 }; 53 54 int offset = co->co_nlocalsplus - co->co_nfreevars; 55 for (int i = 0; i < offset; i++) { 56 update_mapping(i, frame->localsplus[i]); 57 } 58 // Get references to closure variables 59 PyObject* closure = ((PyFunctionObject*)FUNC(frame))->func_closure; 60 for (int i = 0; i < co->co_nfreevars; ++i) { 61 update_mapping(offset + i, PyTuple_GET_ITEM(closure, i)); 62 } 63 64 // NOTE no need to move the instruction pointer to after COPY_FREE_VARS 65 // since we don't actually copy free vars from the closure to the frame 66 // localsplus. 67 68 return mapping.release().ptr(); 69 } 70 71 #endif 72