Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
162 changes: 114 additions & 48 deletions python/_tskitmodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -1786,6 +1786,24 @@ parse_table_collection_dict(tsk_table_collection_t *tables, PyObject *tables_dic
return ret;
}

static const char *
parse_metadata_schema_arg(PyObject *arg, Py_ssize_t* metadata_schema_length)
{
const char *ret = NULL;
if (arg == NULL) {
PyErr_Format(
PyExc_AttributeError,
"Cannot del metadata_schema, set to empty string (\"\") to clear.");
goto out;
}
ret = PyUnicode_AsUTF8AndSize(arg, metadata_schema_length);
if (ret == NULL) {
goto out;
}
out:
return ret;
}

static int
write_table_arrays(tsk_table_collection_t *tables, PyObject *dict)
{
Expand Down Expand Up @@ -2519,16 +2537,10 @@ IndividualTable_set_metadata_schema(IndividualTable *self, PyObject *arg, void *
const char *metadata_schema;
Py_ssize_t metadata_schema_length;

if (arg == NULL) {
PyErr_Format(
PyExc_AttributeError,
"Cannot del metadata_schema, set to empty string (\"\") to clear.");
goto out;
}
if (IndividualTable_check_state(self) != 0) {
goto out;
}
metadata_schema = PyUnicode_AsUTF8AndSize(arg, &metadata_schema_length);
metadata_schema = parse_metadata_schema_arg(arg, &metadata_schema_length);
if (metadata_schema == NULL) {
goto out;
}
Expand Down Expand Up @@ -2995,16 +3007,10 @@ NodeTable_set_metadata_schema(NodeTable *self, PyObject *arg, void *closure)
const char *metadata_schema;
Py_ssize_t metadata_schema_length;

if (arg == NULL) {
PyErr_Format(
PyExc_AttributeError,
"Cannot del metadata_schema, set to empty string (\"\") to clear.");
goto out;
}
if (NodeTable_check_state(self) != 0) {
goto out;
}
metadata_schema = PyUnicode_AsUTF8AndSize(arg, &metadata_schema_length);
metadata_schema = parse_metadata_schema_arg(arg, &metadata_schema_length);
if (metadata_schema == NULL) {
goto out;
}
Expand Down Expand Up @@ -3488,15 +3494,10 @@ EdgeTable_set_metadata_schema(EdgeTable *self, PyObject *arg, void *closure)
const char *metadata_schema;
Py_ssize_t metadata_schema_length;

if (arg == NULL) {
PyErr_Format(PyExc_AttributeError,
"Cannot del metadata_schema, set to empty string (\"\") to clear.");
goto out;
}
if (EdgeTable_check_state(self) != 0) {
goto out;
}
metadata_schema = PyUnicode_AsUTF8AndSize(arg, &metadata_schema_length);
metadata_schema = parse_metadata_schema_arg(arg, &metadata_schema_length);
if (metadata_schema == NULL) {
goto out;
}
Expand Down Expand Up @@ -3989,16 +3990,10 @@ MigrationTable_set_metadata_schema(MigrationTable *self, PyObject *arg, void *cl
const char *metadata_schema;
Py_ssize_t metadata_schema_length;

if (arg == NULL) {
PyErr_Format(
PyExc_AttributeError,
"Cannot del metadata_schema, set to empty string (\"\") to clear.");
goto out;
}
if (MigrationTable_check_state(self) != 0) {
goto out;
}
metadata_schema = PyUnicode_AsUTF8AndSize(arg, &metadata_schema_length);
metadata_schema = parse_metadata_schema_arg(arg, &metadata_schema_length);
if (metadata_schema == NULL) {
goto out;
}
Expand Down Expand Up @@ -4453,16 +4448,10 @@ SiteTable_set_metadata_schema(SiteTable *self, PyObject *arg, void *closure)
const char *metadata_schema;
Py_ssize_t metadata_schema_length;

if (arg == NULL) {
PyErr_Format(
PyExc_AttributeError,
"Cannot del metadata_schema, set to empty string (\"\") to clear.");
goto out;
}
if (SiteTable_check_state(self) != 0) {
goto out;
}
metadata_schema = PyUnicode_AsUTF8AndSize(arg, &metadata_schema_length);
metadata_schema = parse_metadata_schema_arg(arg, &metadata_schema_length);
if (metadata_schema == NULL) {
goto out;
}
Expand Down Expand Up @@ -4956,16 +4945,10 @@ MutationTable_set_metadata_schema(MutationTable *self, PyObject *arg, void *clos
const char *metadata_schema;
Py_ssize_t metadata_schema_length;

if (arg == NULL) {
PyErr_Format(
PyExc_AttributeError,
"Cannot del metadata_schema, set to empty string (\"\") to clear.");
goto out;
}
if (MutationTable_check_state(self) != 0) {
goto out;
}
metadata_schema = PyUnicode_AsUTF8AndSize(arg, &metadata_schema_length);
metadata_schema = parse_metadata_schema_arg(arg, &metadata_schema_length);
if (metadata_schema == NULL) {
goto out;
}
Expand Down Expand Up @@ -5375,16 +5358,10 @@ PopulationTable_set_metadata_schema(PopulationTable *self, PyObject *arg, void *
const char *metadata_schema;
Py_ssize_t metadata_schema_length;

if (arg == NULL) {
PyErr_Format(
PyExc_AttributeError,
"Cannot del metadata_schema, set to empty string (\"\") to clear.");
goto out;
}
if (PopulationTable_check_state(self) != 0) {
goto out;
}
metadata_schema = PyUnicode_AsUTF8AndSize(arg, &metadata_schema_length);
metadata_schema = parse_metadata_schema_arg(arg, &metadata_schema_length);
if (metadata_schema == NULL) {
goto out;
}
Expand Down Expand Up @@ -6083,6 +6060,71 @@ TableCollection_get_file_uuid(TableCollection *self, void *closure)
return Py_BuildValue("s", self->tables->file_uuid);
}

static PyObject *
TableCollection_get_metadata(TableCollection *self, void *closure)
{
return PyBytes_FromStringAndSize(self->tables->metadata, self->tables->metadata_length);
}

static int
TableCollection_set_metadata(TableCollection *self, PyObject *arg, void *closure)
{
int ret = -1;
int err;
char *metadata;
Py_ssize_t metadata_length;

if (arg == NULL) {
PyErr_Format(
PyExc_AttributeError,
"Cannot del metadata, set to empty string (b\"\") to clear.");
goto out;
}
err = PyBytes_AsStringAndSize(arg, &metadata, &metadata_length);
if (err != 0) {
goto out;
}
err = tsk_table_collection_set_metadata(
self->tables, metadata, metadata_length);
if (err != 0) {
handle_library_error(err);
goto out;
}
ret = 0;
out:
return ret;
}

static PyObject *
TableCollection_get_metadata_schema(TableCollection *self, void *closure)
{
return make_Py_Unicode_FromStringAndLength(
self->tables->metadata_schema, self->tables->metadata_schema_length);
}

static int
TableCollection_set_metadata_schema(TableCollection *self, PyObject *arg, void *closure)
{
int ret = -1;
int err;
const char *metadata_schema;
Py_ssize_t metadata_schema_length;

metadata_schema = parse_metadata_schema_arg(arg, &metadata_schema_length);
if (metadata_schema == NULL) {
goto out;
}
err = tsk_table_collection_set_metadata_schema(
self->tables, metadata_schema, metadata_schema_length);
if (err != 0) {
handle_library_error(err);
goto out;
}
ret = 0;
out:
return ret;
}

static PyObject *
TableCollection_simplify(TableCollection *self, PyObject *args, PyObject *kwds)
{
Expand Down Expand Up @@ -6340,6 +6382,12 @@ static PyGetSetDef TableCollection_getsetters[] = {
(setter) TableCollection_set_sequence_length, "The sequence length."},
{"file_uuid", (getter) TableCollection_get_file_uuid, NULL,
"The UUID of the corresponding file."},
{"metadata",
(getter) TableCollection_get_metadata,
(setter) TableCollection_set_metadata, "The metadata."},
{"metadata_schema",
(getter) TableCollection_get_metadata_schema,
(setter) TableCollection_set_metadata_schema, "The metadata schema."},
{NULL} /* Sentinel */
};

Expand Down Expand Up @@ -6680,6 +6728,20 @@ TreeSequence_get_site(TreeSequence *self, PyObject *args)
return ret;
}

static PyObject *
TreeSequence_get_metadata(TreeSequence * self) {
return PyBytes_FromStringAndSize(
self->tree_sequence->tables->metadata,
self->tree_sequence->tables->metadata_length);
}

static PyObject *
TreeSequence_get_metadata_schema(TreeSequence * self) {
return make_Py_Unicode_FromStringAndLength(
self->tree_sequence->tables->metadata_schema,
self->tree_sequence->tables->metadata_schema_length);
}

static PyObject *
TreeSequence_get_table_metadata_schemas(TreeSequence *self) {
PyObject *ret = NULL;
Expand Down Expand Up @@ -8173,6 +8235,10 @@ static PyMethodDef TreeSequence_methods[] = {
METH_NOARGS, "Returns the tree breakpoints as a numpy array." },
{"get_file_uuid", (PyCFunction) TreeSequence_get_file_uuid,
METH_NOARGS, "Returns the UUID of the underlying file, if present." },
{"get_metadata", (PyCFunction) TreeSequence_get_metadata, METH_NOARGS,
"Returns the metadata for the tree sequence"},
{"get_metadata_schema", (PyCFunction) TreeSequence_get_metadata_schema, METH_NOARGS,
"Returns the metadata schema for the tree sequence metadata"},
{"get_num_sites", (PyCFunction) TreeSequence_get_num_sites,
METH_NOARGS, "Returns the number of sites" },
{"get_num_mutations", (PyCFunction) TreeSequence_get_num_mutations, METH_NOARGS,
Expand Down
54 changes: 54 additions & 0 deletions python/tests/test_lowlevel.py
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,38 @@ def test_set_sequence_length(self):
tables.sequence_length = value
self.assertEqual(tables.sequence_length, value)

def test_set_metadata_errors(self):
tables = _tskit.TableCollection(1)
with self.assertRaises(AttributeError):
del tables.metadata
for bad_value in ["bytes only", 59, 43.4, None, []]:
with self.assertRaises(TypeError):
tables.metadata = bad_value

def test_set_metadata(self):
tables = _tskit.TableCollection(1)
self.assertEqual(tables.metadata, b"")
for value in [b"foo", b"", "💩".encode(), b"null char \0 in string"]:
tables.metadata = value
tables.metadata_schema = "Test we have two separate fields"
self.assertEqual(tables.metadata, value)

def test_set_metadata_schema_errors(self):
tables = _tskit.TableCollection(1)
with self.assertRaises(AttributeError):
del tables.metadata_schema
for bad_value in [59, 43.4, None, []]:
with self.assertRaises(TypeError):
tables.metadata_schema = bad_value

def test_set_metadata_schema(self):
tables = _tskit.TableCollection(1)
self.assertEqual(tables.metadata_schema, "")
for value in ["foo", "", "💩", "null char \0 in string"]:
tables.metadata_schema = value
tables.metadata = b"Test we have two separate fields"
self.assertEqual(tables.metadata_schema, value)

def test_simplify_bad_args(self):
ts = msprime.simulate(10, random_seed=1)
tc = ts.tables.ll_tables
Expand Down Expand Up @@ -478,6 +510,28 @@ def test_metadata_schemas(self):
for table_name in self.metadata_tables:
self.assertEqual(getattr(schemas, table_name), "")

def test_metadata(self):
tables = _tskit.TableCollection(1)
ts = _tskit.TreeSequence()
ts.load_tables(tables)
self.assertEqual(ts.get_metadata(), b"")
for value in [b"foo", b"", "💩".encode(), b"null char \0 in string"]:
tables.metadata = value
ts = _tskit.TreeSequence()
ts.load_tables(tables)
self.assertEqual(ts.get_metadata(), value)

def test_metadata_schema(self):
tables = _tskit.TableCollection(1)
ts = _tskit.TreeSequence()
ts.load_tables(tables)
self.assertEqual(ts.get_metadata_schema(), "")
for value in ["foo", "", "💩", "null char \0 in string"]:
tables.metadata_schema = value
ts = _tskit.TreeSequence()
ts.load_tables(tables)
self.assertEqual(ts.get_metadata_schema(), value)

def test_kc_distance_errors(self):
ts1 = self.get_example_tree_sequence(10)
self.assertRaises(TypeError, ts1.get_kc_distance)
Expand Down