Thursday, April 20, 2006

Low-level Programming With Python Extensions

So, Python is a great malleable tool for building up all sorts of things. I use it mainly to build automated tests for a PCI device we're bringing up at work. While the device driver is straight C the tests that use it could be written in pretty much anything. Python is certainly a powerful tool that doesn't get in the way. And everything was groovy until I got around to testing a feature that requires specific structures in PC memory.

For the sake of discussion we'll assume that our device under test wants a pointer to some table structure in host memory and that table then points to several other structures. We write the address of that table to some register in the DUT and say "go." Nothing fancy. Now if we were writing this in C we'd be done by now (except that we prefer not to write these kinds of tests in C for all the usual reasons). In Python however it's not so simple; we can't make a simple struct, and if we could we wouldn't know its address.

After some head scratching you start to remember that it's possible to write extensions for Python in C. You also remembered you tried it once and it scared you. Well let's try it again. Now the simplest thing we could do would be to simply allocate a block of memory and make some functions that let us get at that memory.

Like so:

// A scratch memory module for Python.  Just allocates a block
// of memory and gives the user access to it and its address.

#include <python2.4/Python.h>

#define SCRATCH_SIZE 1024*10 // 10kB

int* scratch_mem_ptr;


PyObject *readw(self, args) PyObject *self, *args; {
PyObject *result = NULL;
int offset;
int val;

if (PyArg_ParseTuple(args, "i", &offset)) {
val = *(scratch_mem_ptr + offset);
result = Py_BuildValue("i", val);
}

return result;
}

PyObject *writew(self, args) PyObject *self, *args; {
PyObject *result = NULL;
int offset;
int val;

if (PyArg_ParseTuple(args, "ii", &offset, &val)) {
*(scratch_mem_ptr + offset) = val;
result = Py_BuildValue("");
}

return result;
}


PyObject *get_addr(self, args) PyObject *self, *args; {
PyObject *result = NULL;

if (PyArg_ParseTuple(args, "")) {
result = Py_BuildValue("i", scratch_mem_ptr);
}

return result;
}

PyMethodDef methods[] = {
{ "readw", readw, METH_VARARGS },
{ "writew", writew, METH_VARARGS },
{ "get_addr", get_addr, METH_VARARGS },
{ NULL, NULL },
};



void initscratchmem() {
(void)Py_InitModule("scratchmem", methods);

scratch_mem_ptr = malloc(SCRATCH_SIZE);
}
scratchmem.c


The above can be compiled like this (the "python2.4" should be swapped out above and below to match your Python version):
gcc -Wall -g -I/usr/include/python2.4/ -c scratchmem.c -o scratchmem.o
ld -shared scratchmem.o -o scratchmem.so

Now let's try it out. To do this you'll need that scratchmem.so to be in a directory Python knows to look in. The simplest is to have it in the current directory.
>>> import scratchmem
>>> hex(scratchmem.get_addr())
'0x81769c8'

Hey, not bad, not bad. Now, let's try to use that memory.
>>> for offset in range(4):
... scratchmem.writew(offset << 2, 0xffff0000 + offset)
...
>>> for offset in range(4):
... print hex(scratchmem.readw(offset << 2))
...
0xffff0000
0xffff0001
0xffff0002
0xffff0003

It works. We're geniuses and we can take the rest of the day off, right? Not quite. We did manage to give our Python programs access to some raw memory but we've knocked ourselves back into the stone ages because we now have to manage that entire block of memory ourselves. A better design would be to have little memory objects that the user can instantiate at will and of varying sizes. This will hopefully be the subject of a future post.

Another thing missing in this design is range checking. This will run amok in someone else's memory if the user gives an offset that overruns what we've malloc'ed, just like our regular C program would.