Support file objects (streams, sockets etc.)

.circleci/config.yml

@@ -114,7 +114,7 @@ jobs:
           name: Run Python tests
           command: |
             cd python
-            python -m pytest --cov=tskit  --cov-report=xml --cov-branch -n `nproc` tests
+            python -m pytest --cov=tskit  --cov-report=xml --cov-branch -n8 tests
 
       - run:
           name: Upload Python coverage

python/CHANGELOG.rst

@@ -8,6 +8,13 @@
   collections with the options to ignore top-level metadata/schema or
   provenance tables. (:user:`mufernando`, :issue:`896`, :pr:`897`).
 
+- ``ts.dump`` and ``tskit.load`` now support reading and writing file objects such as
+  FIFOs and sockets. (:user:`benjeffery`, :issue:`657`, :pr:`909`)
+
+**Breaking changes**
+
+- The argument to ``ts.dump`` and ``tskit.load`` has been renamed `file` from `path`.
+
 --------------------
 [0.3.2] - 2020-09-29
 --------------------

python/_tskitmodule.c

@@ -709,6 +709,33 @@ table_get_column_array(size_t num_rows, void *data, int npy_type, size_t element
     return ret;
 }
 
+static FILE *
+make_file(PyObject *fileobj, const char *mode)
+{
+    FILE *ret = NULL;
+    FILE *file = NULL;
+    int fileobj_fd, new_fd;
+
+    fileobj_fd = PyObject_AsFileDescriptor(fileobj);
+    if (fileobj_fd == -1) {
+        goto out;
+    }
+    new_fd = dup(fileobj_fd);
+    if (new_fd == -1) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        goto out;
+    }
+    file = fdopen(new_fd, mode);
+    if (file == NULL) {
+        (void) close(new_fd);
+        PyErr_SetFromErrno(PyExc_OSError);
+        goto out;
+    }
+    ret = file;
+out:
+    return ret;
+}
+
 /*===================================================================
  * IndividualTable
  *===================================================================
@@ -5321,23 +5348,33 @@ static PyObject *
 TreeSequence_dump(TreeSequence *self, PyObject *args, PyObject *kwds)
 {
     int err;
-    char *path;
+    FILE *file = NULL;
+    PyObject *py_file = NULL;
     PyObject *ret = NULL;
-    static char *kwlist[] = { "path", NULL };
+    static char *kwlist[] = { "file", NULL };
 
     if (TreeSequence_check_tree_sequence(self) != 0) {
         goto out;
     }
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s", kwlist, &path)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O", kwlist, &py_file)) {
+        goto out;
+    }
+
+    file = make_file(py_file, "wb");
+    if (file == NULL) {
         goto out;
     }
-    err = tsk_treeseq_dump(self->tree_sequence, path, 0);
+
+    err = tsk_treeseq_dumpf(self->tree_sequence, file, 0);
     if (err != 0) {
         handle_library_error(err);
         goto out;
     }
     ret = Py_BuildValue("");
 out:
+    if (file != NULL) {
+        (void) fclose(file);
+    }
     return ret;
 }
 
@@ -5398,24 +5435,41 @@ static PyObject *
 TreeSequence_load(TreeSequence *self, PyObject *args, PyObject *kwds)
 {
     int err;
-    char *path;
     PyObject *ret = NULL;
-    static char *kwlist[] = { "path", NULL };
+    PyObject *py_file;
+    FILE *file = NULL;
+    static char *kwlist[] = { "file", NULL };
 
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s", kwlist, &path)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O", kwlist, &py_file)) {
+        goto out;
+    }
+    file = make_file(py_file, "rb");
+    if (file == NULL) {
+        goto out;
+    }
+    /* Set unbuffered mode to ensure no more bytes are read than requested.
+     * Buffered reads could read beyond the end of the current store in a
+     * multi-store file or stream. This data would be discarded when we
+     * fclose() the file below, such that attempts to load the next store
+     * will fail. */
+    if (setvbuf(file, NULL, _IONBF, 0) != 0) {
+        PyErr_SetFromErrno(PyExc_OSError);
         goto out;
     }
     err = TreeSequence_alloc(self);
     if (err != 0) {
         goto out;
     }
-    err = tsk_treeseq_load(self->tree_sequence, path, 0);
+    err = tsk_treeseq_loadf(self->tree_sequence, file, 0);
     if (err != 0) {
         handle_library_error(err);
         goto out;
     }
     ret = Py_BuildValue("");
 out:
+    if (file != NULL) {
+        (void) fclose(file);
+    }
     return ret;
 }
 

python/tests/conftest.py

@@ -1,13 +1,62 @@
+# MIT License
+#
+# Copyright (c) 2018-2020 Tskit Developers
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""
+Configuration and fixtures for pytest. Only put test-suite wide fixtures in here. Module
+specific fixtures should live in their modules.
+
+To use a fixture in a test simply refer to it by name as an argument. This is called
+dependancy injection. Note that all fixtures should have the suffix "_fixture" to make
+it clear in test code.
+
+For example to use the `ts` fixture (a tree sequence with data in all fields) in a test:
+
+def test_something(ts):
+    assert ts.some_method() == expected
+
+Fixtures can be parameterised etc. see https://docs.pytest.org/en/stable/fixture.html
+
+Note that fixtures have a "scope" for example `ts` below is only created once per
+test session and re-used for subsequent tests.
+"""
+import msprime
 import pytest
+from pytest import fixture
+
+import tskit
 
 
 def pytest_addoption(parser):
+    """
+    Add an option to skip tests marked with `@pytest.mark.slow`
+    """
     parser.addoption(
         "--skip-slow", action="store_true", default=False, help="Skip slow tests"
     )
 
 
 def pytest_configure(config):
+    """
+    Add docs on the "slow" marker
+    """
     config.addinivalue_line("markers", "slow: mark test as slow to run")
 
 
@@ -17,3 +66,64 @@ def pytest_collection_modifyitems(config, items):
         for item in items:
             if "slow" in item.keywords:
                 item.add_marker(skip_slow)
+
+
+@fixture(scope="session")
+def simple_ts_fixture():
+    return msprime.simulate(2, random_seed=42)
+
+
+@fixture(scope="session")
+def ts_fixture():
+    """
+    A tree sequence with data in all fields
+    """
+    n = 10
+    t = 1
+    population_configurations = [
+        msprime.PopulationConfiguration(n // 2),
+        msprime.PopulationConfiguration(n // 2),
+        msprime.PopulationConfiguration(0),
+    ]
+    demographic_events = [
+        msprime.MassMigration(time=t, source=0, destination=2),
+        msprime.MassMigration(time=t, source=1, destination=2),
+    ]
+    ts = msprime.simulate(
+        population_configurations=population_configurations,
+        demographic_events=demographic_events,
+        random_seed=1,
+        mutation_rate=1,
+        record_migrations=True,
+    )
+    tables = ts.dump_tables()
+    for table in [
+        "edges",
+        "individuals",
+        "migrations",
+        "mutations",
+        "nodes",
+        "populations",
+        "sites",
+    ]:
+        getattr(tables, table).metadata_schema = tskit.MetadataSchema({"codec": "json"})
+        metadatas = [f"n_{table}_{u}" for u in range(getattr(ts, f"num_{table}"))]
+        metadata, metadata_offset = tskit.pack_strings(metadatas)
+        getattr(tables, table).set_columns(
+            **{
+                **getattr(tables, table).asdict(),
+                "metadata": metadata,
+                "metadata_offset": metadata_offset,
+            }
+        )
+    tables.metadata_schema = tskit.MetadataSchema({"codec": "json"})
+    tables.metadata = "Test metadata"
+    return tables.tree_sequence()
+
+
+@fixture(scope="session")
+def replicate_ts_fixture():
+    """
+    A list of tree sequences
+    """
+    return list(msprime.simulate(10, num_replicates=10, random_seed=42))

python/tests/test_cli.py

@@ -556,7 +556,7 @@ def verify(self, command):
             with pytest.raises(TestException):
                 capture_output(cli.tskit_main, ["info", "/no/such/file"])
             mocked_exit.assert_called_once_with(
-                "Load error: [Errno 2] No such file or directory"
+                "Load error: [Errno 2] No such file or directory: '/no/such/file'"
             )
 
     def test_info(self):