Merge pull request #1832 from benjeffery/time-units-enumeration

mergify[bot] · web-flow · commit cc9e445d560b · 2021-10-22T14:05:38.000Z
Add time_units enumeration
diff --git a/c/tests/test_tables.c b/c/tests/test_tables.c
@@ -454,7 +454,7 @@ test_table_collection_time_units(void)
     tsk_table_collection_free(&tc2);
     ret = tsk_table_collection_init(&tc1, 0);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
-    CU_ASSERT_EQUAL_FATAL(ret, strncmp(tc1.time_units, TSK_DEFAULT_TIME_UNITS, 7));
+    CU_ASSERT_EQUAL_FATAL(ret, strncmp(tc1.time_units, TSK_TIME_UNITS_UNKNOWN, 7));
     tc1.sequence_length = 1.0;
     ret = tsk_table_collection_dump(&tc1, _tmp_file_name, 0);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
diff --git a/c/tskit/core.h b/c/tskit/core.h
@@ -87,7 +87,7 @@ __tsk_nan_f(void)
 }
 #define TSK_UNKNOWN_TIME __tsk_nan_f()
 
-#define TSK_DEFAULT_TIME_UNITS "unknown"
+#define TSK_TIME_UNITS_UNKNOWN "unknown"
 #define TSK_TIME_UNITS_UNCALIBRATED "uncalibrated"
 
 /**
diff --git a/c/tskit/tables.c b/c/tskit/tables.c
@@ -9789,7 +9789,7 @@ tsk_table_collection_init(tsk_table_collection_t *self, tsk_flags_t options)
 
     /* Set default time_units value */
     ret = tsk_table_collection_set_time_units(
-        self, TSK_DEFAULT_TIME_UNITS, strlen(TSK_DEFAULT_TIME_UNITS));
+        self, TSK_TIME_UNITS_UNKNOWN, strlen(TSK_TIME_UNITS_UNKNOWN));
     if (ret != 0) {
         goto out;
     }
diff --git a/python/CHANGELOG.rst b/python/CHANGELOG.rst
@@ -56,7 +56,7 @@
 
 - Add ``time_units`` to ``TreeSequence`` to describe the units of the time dimension of the
   tree sequence. This is then used to generate an error if ``time_units`` is ``uncalibrated`` when
-  using the branch lengths in statistics. (:user:`benjeffery`, :issue:`1644`, :pr:`1760`)
+  using the branch lengths in statistics. (:user:`benjeffery`, :issue:`1644`, :pr:`1760`, :pr:`1832`)
 
 - Add the ``virtual_root`` property to the Tree class (:user:`jeromekelleher`, :pr:`1704`).
 
diff --git a/python/_tskitmodule.c b/python/_tskitmodule.c
@@ -11915,5 +11915,9 @@ PyInit__tskit(void)
     PyModule_AddIntConstant(module, "FORWARD", TSK_DIR_FORWARD);
     PyModule_AddIntConstant(module, "REVERSE", TSK_DIR_REVERSE);
 
+    PyModule_AddStringConstant(module, "TIME_UNITS_UNKNOWN", TSK_TIME_UNITS_UNKNOWN);
+    PyModule_AddStringConstant(
+        module, "TIME_UNITS_UNCALIBRATED", TSK_TIME_UNITS_UNCALIBRATED);
+
     return module;
 }
diff --git a/python/lwt_interface/dict_encoding_testlib.py b/python/lwt_interface/dict_encoding_testlib.py
@@ -211,7 +211,7 @@ def test_missing_time_units(self, tables):
         lwt = lwt_module.LightweightTableCollection()
         lwt.fromdict(d)
         tables = tskit.TableCollection.fromdict(lwt.asdict())
-        assert tables.time_units == "unknown"
+        assert tables.time_units == tskit.TIME_UNITS_UNKNOWN
 
     def test_missing_metadata(self, tables):
         assert tables.metadata != b""
@@ -603,7 +603,7 @@ def test_top_level_time_units(self, tables):
         lwt.fromdict(d)
         out = lwt.asdict()
         tables = tskit.TableCollection.fromdict(out)
-        assert tables.time_units == "unknown"
+        assert tables.time_units == tskit.TIME_UNITS_UNKNOWN
         # Missing is tested in TestMissingData above
         d = tables.asdict()
         # None should give default value
diff --git a/python/tests/test_highlevel.py b/python/tests/test_highlevel.py
@@ -2402,14 +2402,16 @@ def test_tree_sequence_metadata(self):
     def test_tree_sequence_time_units(self):
         tc = tskit.TableCollection(1)
         ts = tc.tree_sequence()
-        assert ts.time_units == "unknown"
+        assert ts.time_units == tskit.TIME_UNITS_UNKNOWN
         tc.time_units = "something else"
         ts = tc.tree_sequence()
         assert ts.time_units == "something else"
         with pytest.raises(AttributeError):
             del ts.time_units
         with pytest.raises(AttributeError):
             ts.time_units = "readonly"
+        assert tskit.TIME_UNITS_UNKNOWN == "unknown"
+        assert tskit.TIME_UNITS_UNCALIBRATED == "uncalibrated"
 
     def test_table_metadata_schemas(self):
         ts = msprime.simulate(5)
diff --git a/python/tests/test_lowlevel.py b/python/tests/test_lowlevel.py
@@ -234,7 +234,7 @@ def test_set_time_units_errors(self):
 
     def test_set_time_units(self):
         tables = _tskit.TableCollection(1)
-        assert tables.time_units == "unknown"
+        assert tables.time_units == tskit.TIME_UNITS_UNKNOWN
         for value in ["foo", "", "💩", "null char \0 in string"]:
             tables.time_units = value
             assert tables.time_units == value
@@ -1309,7 +1309,7 @@ def test_time_units(self):
         tables.build_index()
         ts = _tskit.TreeSequence()
         ts.load_tables(tables)
-        assert ts.get_time_units() == "unknown"
+        assert ts.get_time_units() == tskit.TIME_UNITS_UNKNOWN
         for value in ["foo", "", "💩", "null char \0 in string"]:
             tables.time_units = value
             ts = _tskit.TreeSequence()
@@ -3254,3 +3254,8 @@ def test_uninitialised():
                 method = getattr(uninitialised, method_name)
                 with pytest.raises((SystemError, ValueError)):
                     method()
+
+
+def test_constants():
+    assert _tskit.TIME_UNITS_UNKNOWN == "unknown"
+    assert _tskit.TIME_UNITS_UNCALIBRATED == "uncalibrated"
diff --git a/python/tests/test_tables.py b/python/tests/test_tables.py
@@ -4078,7 +4078,7 @@ def test_set_metadata(self):
 
     def test_set_time_units(self):
         tc = tskit.TableCollection(1)
-        assert tc.time_units == "unknown"
+        assert tc.time_units == tskit.TIME_UNITS_UNKNOWN
 
         ex1 = "years"
         ex2 = "generations"
diff --git a/python/tests/test_tree_stats.py b/python/tests/test_tree_stats.py
@@ -6084,7 +6084,7 @@ class TestTimeUncalibratedErrors:
     def test_uncalibrated_time_allele_frequency_spectrum(self, ts_fixture):
         ts_fixture.allele_frequency_spectrum(mode="branch")
         tables = ts_fixture.dump_tables()
-        tables.time_units = "uncalibrated"
+        tables.time_units = tskit.TIME_UNITS_UNCALIBRATED
         ts_uncalibrated = tables.tree_sequence()
         ts_uncalibrated.allele_frequency_spectrum(mode="site")
         with pytest.raises(
@@ -6100,7 +6100,7 @@ def test_uncalibrated_time_general_stat(self, ts_fixture):
             W, lambda x: x * (x < ts_fixture.num_samples), W.shape[1], mode="branch"
         )
         tables = ts_fixture.dump_tables()
-        tables.time_units = "uncalibrated"
+        tables.time_units = tskit.TIME_UNITS_UNCALIBRATED
         ts_uncalibrated = tables.tree_sequence()
         ts_uncalibrated.general_stat(
             W, lambda x: x * (x < ts_uncalibrated.num_samples), W.shape[1], mode="site"
diff --git a/python/tskit/__init__.py b/python/tskit/__init__.py
@@ -50,6 +50,12 @@
 #: NAN value, you cannot use `==` to test for it. Use :func:`is_unknown_time` instead.
 UNKNOWN_TIME = _tskit.UNKNOWN_TIME
 
+#: Default value of ts.time_units
+TIME_UNITS_UNKNOWN = _tskit.TIME_UNITS_UNKNOWN
+
+#: ts.time_units value when dimension is uncalibrated
+TIME_UNITS_UNCALIBRATED = _tskit.TIME_UNITS_UNCALIBRATED
+
 #: Options for printing to strings and HTML, modify with tskit.set_print_options.
 _print_options = {"max_lines": 40}
 
diff --git a/python/tskit/drawing.py b/python/tskit/drawing.py
@@ -38,6 +38,7 @@
 import numpy as np
 import svgwrite
 
+import tskit
 import tskit.util as util
 from _tskit import NODE_IS_SAMPLE
 from _tskit import NULL
@@ -551,7 +552,7 @@ def __init__(
                 y_label = "Node time"
             else:
                 y_label = "Time"
-            if ts.time_units != "unknown":
+            if ts.time_units != tskit.TIME_UNITS_UNKNOWN:
                 y_label += f" ({ts.time_units})"
         self.x_label = x_label
         self.y_label = y_label
diff --git a/python/tskit/trees.py b/python/tskit/trees.py
@@ -41,6 +41,7 @@
 import numpy as np
 
 import _tskit
+import tskit
 import tskit.combinatorics as combinatorics
 import tskit.drawing as drawing
 import tskit.exceptions as exceptions
@@ -3153,7 +3154,7 @@ def parse_mutations(
         if len(tokens) >= 3:
             site = int(tokens[site_index])
             node = int(tokens[node_index])
-            if time_index is None or tokens[time_index] == "unknown":
+            if time_index is None or tokens[time_index] == tskit.TIME_UNITS_UNKNOWN:
                 time = UNKNOWN_TIME
             else:
                 time = float(tokens[time_index])

Original file line number	Diff line number	Diff line change
`@@ -87,7 +87,7 @@ __tsk_nan_f(void)`
`87`	`87`	`}`
`88`	`88`	`#define TSK_UNKNOWN_TIME __tsk_nan_f()`
`89`	`89`
`90`		`-#define TSK_DEFAULT_TIME_UNITS "unknown"`
	`90`	`+#define TSK_TIME_UNITS_UNKNOWN "unknown"`
`91`	`91`	`#define TSK_TIME_UNITS_UNCALIBRATED "uncalibrated"`
`92`	`92`
`93`	`93`	`/**`
Original file line number	Diff line number	Diff line change
`@@ -9789,7 +9789,7 @@ tsk_table_collection_init(tsk_table_collection_t *self, tsk_flags_t options)`
`9789`	`9789`
`9790`	`9790`	`/* Set default time_units value */`
`9791`	`9791`	`ret = tsk_table_collection_set_time_units(`
`9792`		`- self, TSK_DEFAULT_TIME_UNITS, strlen(TSK_DEFAULT_TIME_UNITS));`
	`9792`	`+ self, TSK_TIME_UNITS_UNKNOWN, strlen(TSK_TIME_UNITS_UNKNOWN));`
`9793`	`9793`	`if (ret != 0) {`
`9794`	`9794`	`goto out;`
`9795`	`9795`	`}`
Original file line number	Diff line number	Diff line change
`@@ -11915,5 +11915,9 @@ PyInit__tskit(void)`
`11915`	`11915`	`PyModule_AddIntConstant(module, "FORWARD", TSK_DIR_FORWARD);`
`11916`	`11916`	`PyModule_AddIntConstant(module, "REVERSE", TSK_DIR_REVERSE);`
`11917`	`11917`
	`11918`	`+ PyModule_AddStringConstant(module, "TIME_UNITS_UNKNOWN", TSK_TIME_UNITS_UNKNOWN);`
	`11919`	`+ PyModule_AddStringConstant(`
	`11920`	`+ module, "TIME_UNITS_UNCALIBRATED", TSK_TIME_UNITS_UNCALIBRATED);`
	`11921`	`+`
`11918`	`11922`	`return module;`
`11919`	`11923`	`}`