Writing C Extensions for Python

A Minimal Working Example

The vast majority of my code these days is written in Python, and recently I hit a point where it just started to feel painfully slow. Every other day I was thinking to myself, “there has got to be a way to make this run faster”. There are multiple strategies to achieve this (and since I’ve been wanting to learn C for a while), so I took this opportunity to make a C extension in Python.

When writing a C extension, your C code will need the following 4 pieces:

1. A function returning a PyObject* value — the code that can be called in native Python
2. A PyMethodDef to define the mapping between Python method names and the C functions
3. A PyModuleDef to define the module that contains your methods
4. An initialization function named PyInit_[module] to initialize your module when it’s imported

With this, you can create functions written in C that can then be compiled and called by native Python code. Let’s go through an example.

Create your function

Let’s create a simple C function that will return a PyObject* struct — in this case, a PyLong_FromLong, a.k.a. an integer:

// example.c

#include <stdio.h>
#include <stdlib.h>
#include <Python.h>

static PyObject *method_count_elements(PyObject *self, PyObject *args) {
  PyObject *py_list;

  if (!PyArg_ParseTuple(args, "O", &py_list)) {
    // returning null raises an exception to the Python caller
    return NULL; 
  }

  if (!PyList_Check(py_list)) {
    PyErr_SetString(PyExc_TypeError, "Expected a list");
    return NULL;
  }
  Py_ssize_t len = PyList_Size(py_list);
  
  return PyLong_FromLong(len);
}

Create the PyMethodDef struct

static PyMethodDef CustomMethods[] = {
  {"count_elements", method_count_elements, METH_VARARGS, "Counts the elements in a list"},
  {NULL, NULL, 0, NULL},
};

Create the PyModuleDef struct

static struct PyModuleDef custom_extension_module = {
  PyModuleDef_HEAD_INIT,
  "custom_extension",
  "Python interface for custom C extension",
  -1,
  CustomMethods,
};

Create the Initialization Function

PyMODINIT_FUNC PyInit_custom_extension(void) {
    return PyModule_Create(&custom_extension_module);
}

Once these are all together, you can create a setup.py file to compile the code as a Python module.

import os
import sysconfig
from setuptools import Extension, setup


# depending on your virtual environment, you may need to specify
# the path to the Python headers (Python.h) or library directories
# for linking
python_header_path = sysconfig.get_paths()["include"]
include_dirs = [python_header_path]

# it may also be necessary to specify library directories so you can
# find required .a or .dylib files to compile. For example, if you 
# want to make SSH connections with libssh2, you may need to specify
# which directories to search to find the libssh2.a and libssh2.dylib
# files
library_dirs = []

# libraries correspond to flags that you would pass to a compiler to 
# ensure that certain libraries are linked to your code at compile time.
# For example, if you wanted to link libssh2 using the compiler, you 
# would pass the flag -lssh2. With the libraries option, you can 
# simply set libraries equal to ["ssh2"]
libraries = []

setup(
  name="custom-extension",     # package name
  version="0.0.1",             # package version
  ext_modules=[
    Extension(
      name="custom_extension",           
      sources=["example.c"],  
      include_dirs=include_dirs,         
      library_dirs=library_dirs,        
      libraries=libraries,
      extra_compile_args=["-Wall", "-O2"],
    )
  ],
)

When you define the setup portion of the setup.py file, you can set the name and the version for your package. When you run pip list, this name and version is what will be displayed once the package is installed. The name you define under Extension is the name that you will use to import the module into your Python files. It is also used to determine where to place the compiled shared object library (.so file).

Get Tate Wilks’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

If you have other C or C header files that you want to include, you can specify which directory to look for them in the include_dirs option.

Now that you have all the key pieces together, you can run pip install . to install the module in your virtual environment. See that it’s installed in your environment by running pip list | grep custom_extension. You can then import the module into your Python files and call the underlying C function:

import custom_extension

if __name__ == "__main__":
    lst = [chr(i) for i in range(97, 123)]
    num_elements = custom_extension.count_elements(lst)
    print(f"Counted {num_elements} element(s) in list.")

From here, you can create whatever functions you want that leverage optimized C code. A recent use case where I applied this was executing commands via SSH to verify the number of services running on a given server. I leveraged the pthread library to achieve concurrency and the resulting read execution took an average of 0.43 seconds for multiple hosts, compared to some executions that took up to 100 seconds.

Bonus: Add a .pyi file

It can be helpful to define a pyi file (Python interface stub file) which provides your IDE with type information about your C extension. A pyi file typically includes the function prototype with docstring information for your IDE to refer to. It won’t affect runtime behavior, but it improves autocompletion and static analysis. In the same location as your source C file, you can add the following pyi file:

def count_elements(lst: list) -> int:
  """  
  Count the number of elements in a list.

  Args:
    lst (list): A Python list to evaluate

  Returns:
    int: The number of elements in the list
"""

You can download all source files for this example — including the C source, setup script, and stub file — from my Gist.