data-apis / python-record-api

Inferring Python API signatures from tracing usage.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Cleanup of Python Stack Access

saulshanabrook opened this issue · comments

There were some comments here by @Caagr98 on how we can clean up our ability to access the Python stack, that we could try to incorporate. I haven't looked into them yet.

Here's the simplified code I ended up with.

import sys
import ctypes as c

__all__ = ("get_stack",)

class Frame(c.Structure):
	_fields_ = (
		*(
			("_ob_next", c.POINTER(c.py_object)),
			("_ob_prev", c.POINTER(c.py_object)),
		) * sys.flags.debug,
		("ob_refcnt", c.c_ssize_t),
		("ob_type", c.py_object),
		("ob_size", c.c_ssize_t),
		("f_back", c.py_object), # c.POINTER(Frame)
		("f_code", c.py_object),
		("f_builtins", c.py_object),
		("f_globals", c.py_object),
		("f_locals", c.py_object),
		# The two fields below are pointers to PyObject arrays, but ctypes is
		# kinda weird so it's easier to just use them as void pointers
		("f_valuestack", c.c_void_p),
		("f_stacktop", c.c_void_p),
	)

def get_stack(frame):
	PTR_SIZE = c.sizeof(c.POINTER(c.py_object))
	_frame = Frame.from_address(id(frame))

	return tuple(
		c.py_object.from_address(addr).value
		for addr in range(_frame.f_valuestack, _frame.f_stacktop, PTR_SIZE)
	)

I'm not 100% certain I don't need to do anything about refcounts, but py_object probably deals with that automatically.

@Caagr98 Awesome, thank you for this!

If anyone feels like updating our stack functionality to use this logic, you can edit the record_api/get_stack.py file.

In our usage we just need to be able to look at the ith item from the top of the stack, we never need to iterate through it all. So all we really need is the __getitem__ function to work on that OpStack class with negative indexes. Since performance is a concern here, we don't have any bounds checks.