diff --git a/python/_tskitmodule.c b/python/_tskitmodule.c index 1d2ce81500..2a1a987785 100644 --- a/python/_tskitmodule.c +++ b/python/_tskitmodule.c @@ -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) { @@ -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; } @@ -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; } @@ -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; } @@ -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; } @@ -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; } @@ -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; } @@ -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; } @@ -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) { @@ -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 */ }; @@ -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; @@ -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, diff --git a/python/tests/test_lowlevel.py b/python/tests/test_lowlevel.py index 8cdcad44a9..ed65d7233d 100644 --- a/python/tests/test_lowlevel.py +++ b/python/tests/test_lowlevel.py @@ -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 @@ -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)