Flesh out examples and add docs to LWT

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.
+
+
+
+
+
+
+
+

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;
 }

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