From cf31f4f19d97471188433a77169911d59ed6f894 Mon Sep 17 00:00:00 2001 From: Ben Jeffery Date: Tue, 17 Nov 2020 00:46:06 +0000 Subject: [PATCH] Flesh out examples and add docs to LWT --- python/lwt_interface/README.md | 47 +++++++++- python/lwt_interface/example_c_module.c | 91 +++++++++++++++++++ python/lwt_interface/test_example_c_module.py | 38 ++++++++ 3 files changed, 175 insertions(+), 1 deletion(-) diff --git a/python/lwt_interface/README.md b/python/lwt_interface/README.md index 09f860b2b4..6e71789be6 100644 --- a/python/lwt_interface/README.md +++ b/python/lwt_interface/README.md @@ -8,4 +8,49 @@ C API in your own compiled Python module (either via Cython or the Python C API), you almost certainly don't need to use this code. -## TODO document +## Overview + +To allow a tskit table collection to be transferred from one compiled Python +extension module to another the table collection is converted to a `dict` of +basic python types and numpy arrays. This is then converted back in the receiving +module. `tskit_lwt_interface.h` provides a function `register_lwt_class` that +defines a Python class `LightweightTableCollection` that performs these conversions +with methods `asdict` and `fromdict`. These methods mirror the `asdict` and `fromdict` +methods on `tskit.TableCollection`. + +## Usage +An example C module skeleton `example_c_module.c` is provided, which shows passing tables +to the C module. See `test_example_c_module.py` for the python example usage +of the example module. + +To add the +`LightweightTableCollection` type to your module you include `tskit_lwt_interface.h` +and then call `register_lwt_class` on your C Python module object. You can then convert +to and from the lightweight table collection in Python, for example to convert a tskit +`TableCollection` to a `LightweightTableCollection`: +```python +tables = tskit.TableCollection(1) +lwt = example_c_module.LightweightTableCollection() +lwt.fromdict(tables.asdict()) +``` +and vice-versa: +```python +tc = tskit.TableCollection(lwt.asdict()) +``` +In C you can access the tables in a `LightweightTableCollection` instance that is passed +to your function, as shown in the `example_receiving` function in `example_c_module.c`. +Note the requirement to check for errors from tskit functions and to call +`handle_tskit_error` to set a Python error, returning `NULL` to Python to indicate error. + +Tables can also be modified in the extension code as in `example_modifying`. We recommend +creating table collections in Python then passing them to C for modification rather than +creating them in C and returning them. This avoids complex managing of object lifecycles +in C code. + + + + + + + + diff --git a/python/lwt_interface/example_c_module.c b/python/lwt_interface/example_c_module.c index 3e2bd98c70..f7849f8bac 100644 --- a/python/lwt_interface/example_c_module.c +++ b/python/lwt_interface/example_c_module.c @@ -38,7 +38,96 @@ #include "tskit_lwt_interface.h" +static PyObject * +example_receiving(PyObject *self, PyObject *args) { + int err = -1; + PyObject* ret = NULL; + LightweightTableCollection *tables = NULL; + tsk_treeseq_t tree_seq; + tsk_tree_t tree; + + memset(&tree, 0, sizeof(tsk_tree_t)); + memset(&tree_seq, 0, sizeof(tsk_treeseq_t)); + + /* Get the tables from the args */ + if (!PyArg_ParseTuple(args, "O!", &LightweightTableCollectionType, &tables)) { + goto out; + } + + /* Check that the tables are init'd to prevent seg faults */ + if (LightweightTableCollection_check_state(tables) != 0) { + goto out; + } + + /* Build a tree sequence from the tables */ + err = tsk_treeseq_init(&tree_seq, tables->tables, 0); + if (err < 0) { + handle_tskit_error(err); + goto out; + } + + /* Get the first tree */ + err = tsk_tree_init(&tree, &tree_seq, 0); + if (err < 0) { + handle_tskit_error(err); + goto out; + } + err = tsk_tree_first(&tree); + if (err < 0) { + handle_tskit_error(err); + goto out; + } + + /* Return true if the tree has more than one root */ + ret = Py_BuildValue("O", tsk_tree_get_num_roots(&tree) > 1 ? Py_True: Py_False); + +out: + tsk_tree_free(&tree); + tsk_treeseq_free(&tree_seq); + return ret; +} + +static PyObject * example_modifying(PyObject *self, PyObject *args) { + int err = -1; + PyObject* ret = NULL; + LightweightTableCollection *tables = NULL; + + if (!PyArg_ParseTuple(args, "O!", &LightweightTableCollectionType, &tables)) { + goto out; + } + + /* Check that the tables are init'd to prevent seg faults */ + if (LightweightTableCollection_check_state(tables) != 0) { + goto out; + } + + /* Modify the tables, note the need to check for error states and handle them */ + err = tsk_table_collection_clear(tables->tables, 0); + if (err < 0) { + handle_tskit_error(err); + goto out; + } + err = tsk_node_table_add_row(&tables->tables->nodes, 0, 0, 0, 0, NULL, 0); + if (err < 0) { + handle_tskit_error(err); + goto out; + } + err = tsk_node_table_add_row(&tables->tables->nodes, 0, 0, 0, 0, NULL, 0); + if (err < 0) { + handle_tskit_error(err); + goto out; + } + + /* Only set the return after no errors */ + ret = Py_BuildValue(""); +out: + return ret; +} + + static PyMethodDef example_c_module_methods[] = { + {"example_receiving", (PyCFunction) example_receiving, METH_VARARGS, "Example of function receiving tables"}, + {"example_modifying", (PyCFunction) example_modifying, METH_VARARGS, "Example of function modifying tables"}, { NULL, NULL, 0, NULL } /* sentinel */ }; @@ -60,7 +149,9 @@ PyInit_example_c_module(void) if (register_lwt_class(module) != 0) { return NULL; } + /* Put your own functions/class definitions here, as usual */ + return module; } diff --git a/python/lwt_interface/test_example_c_module.py b/python/lwt_interface/test_example_c_module.py index 5a308b0e14..20cfd1c255 100644 --- a/python/lwt_interface/test_example_c_module.py +++ b/python/lwt_interface/test_example_c_module.py @@ -2,13 +2,17 @@ import os import sys +import pytest + # Make sure we use the local tskit version. + sys.path.insert(0, os.path.abspath("../")) # An example of how to run the tests defined in the dict_encoding_testlib.py # file for a given compiled version of the code. import dict_encoding_testlib import example_c_module +import tskit # The test cases defined in dict_encoding_testlib all use the form # lwt_module.LightweightTableCollection() to create an instance @@ -18,3 +22,37 @@ dict_encoding_testlib.lwt_module = example_c_module from dict_encoding_testlib import * + + +def test_example_receiving(): + # The example_receiving function returns true if the first tree + # has more than one root + lwt = example_c_module.LightweightTableCollection() + tables = tskit.TableCollection(1) + lwt.fromdict(tables.asdict()) + # Our example function throws an error for an empty table collection + with pytest.raises(ValueError, match="Table collection must be indexed"): + example_c_module.example_receiving(lwt) + + # This tree sequence has one root so we get false + tables = msprime.simulate(10).tables + lwt.fromdict(tables.asdict()) + assert not example_c_module.example_receiving(lwt) + + # Add a root and we get true + tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE) + lwt.fromdict(tables.asdict()) + assert example_c_module.example_receiving(lwt) + + +def test_example_modifying(): + lwt = example_c_module.LightweightTableCollection() + # The example_modifying function clears out the table and adds two rows + tables = msprime.simulate(10, random_seed=42).tables + assert tables.edges.num_rows == 18 + assert tables.nodes.num_rows == 19 + lwt.fromdict(tables.asdict()) + example_c_module.example_modifying(lwt) + modified_tables = tskit.TableCollection.fromdict(lwt.asdict()) + assert modified_tables.edges.num_rows == 0 + assert modified_tables.nodes.num_rows == 2