.

pyslim/methods.py

@@ -55,7 +55,8 @@ def recapitate(ts,
     '''
     is_current_version(ts, _warn=True)
     has_null = False
-    for n in ts.samples():
+    for nid in ts.samples():
+        n = ts.node(nid)
         if n.metadata['is_null']:
             has_null = True
             break

tests/test_annotation.py

@@ -6,6 +6,7 @@
 import pytest
 import random
 import warnings
+import contextlib
 
 import msprime
 import tskit
@@ -14,6 +15,13 @@
 import tests
 from .recipe_specs import restarted_recipe_eq
 
+def mutcontext(ts):
+    if ts.num_mutations > 0:
+        handler = pytest.warns(Warning, match="already has.*metadata")
+    else:
+        handler = contextlib.nullcontext()
+    return handler
+
 class TestAnnotate(tests.PyslimTestCase):
     '''
     Tests for tools to annotate existing msprime-derived tree sequences.
@@ -270,13 +278,18 @@ def test_basic_annotation(self, helper_functions, tmp_path):
                 tick = 4
                 cycle = 1
                 stage = "late"
-                slim_ts = pyslim.annotate(
-                        ts, model_type="WF",
-                        tick=tick,
-                        cycle=cycle,
-                        stage=stage,
-                        annotate_mutations=do_mutations,
-                )
+                if do_mutations:
+                    handler = mutcontext(ts)
+                else:
+                    handler = contextlib.nullcontext()
+                with handler:
+                    slim_ts = pyslim.annotate(
+                            ts, model_type="WF",
+                            tick=tick,
+                            cycle=cycle,
+                            stage=stage,
+                            annotate_mutations=do_mutations,
+                    )
                 assert slim_ts.metadata['SLiM']['model_type'] == 'WF'
                 assert slim_ts.metadata['SLiM']['tick'] == tick
                 assert slim_ts.metadata['SLiM']['cycle'] == cycle
@@ -305,7 +318,8 @@ def test_annotate_refseq(self):
     def test_annotate_individuals(self, helper_functions, tmp_path):
         # test the workflow of annotating defaults, then assigning sexes randomly
         for ts in helper_functions.get_msprime_examples():
-            slim_ts = pyslim.annotate(ts, model_type="nonWF", tick=1, stage="early")
+            with mutcontext(ts):
+                slim_ts = pyslim.annotate(ts, model_type="nonWF", tick=1, stage="early")
             tables = slim_ts.dump_tables()
             top_md = tables.metadata
             top_md['SLiM']['separate_sexes'] = True
@@ -345,7 +359,8 @@ def test_annotate_XY(self, helper_functions, tmp_path):
         random.seed(8)
         for ts in helper_functions.get_msprime_examples():
             for genome_type in ["X", "Y"]:
-                slim_ts = pyslim.annotate_defaults(ts, model_type="nonWF", slim_generation=1)
+                with mutcontext(ts):
+                    slim_ts = pyslim.annotate_defaults(ts, model_type="nonWF", slim_generation=1)
                 tables = slim_ts.dump_tables()
                 top_md = tables.metadata
                 top_md['SLiM']['separate_sexes'] = True
@@ -399,7 +414,8 @@ def test_annotate_XY(self, helper_functions, tmp_path):
     def test_annotate_nodes(self, helper_functions):
         # test workflow of annotating defaults and then editing node metadata
         for ts in helper_functions.get_msprime_examples():
-            slim_ts = pyslim.annotate(ts, model_type="nonWF", tick=1)
+            with mutcontext(ts):
+                slim_ts = pyslim.annotate(ts, model_type="nonWF", tick=1)
             tables = slim_ts.dump_tables()
             metadata = [n.metadata for n in tables.nodes]
             gtypes = [
@@ -421,7 +437,8 @@ def test_annotate_nodes(self, helper_functions):
 
     def test_annotate_mutations(self, helper_functions):
         for ts in helper_functions.get_msprime_examples():
-            slim_ts = pyslim.annotate(ts, model_type="nonWF", tick=1)
+            with mutcontext(ts):
+                slim_ts = pyslim.annotate(ts, model_type="nonWF", tick=1)
             tables = slim_ts.dump_tables()
             metadata = [m.metadata for m in tables.mutations]
             selcoefs = [random.uniform(0, 1) for _ in metadata]

tests/test_tree_sequence.py

@@ -637,7 +637,7 @@ def test_mutation_at(self, recipe):
         ts = recipe["ts"]
         for _ in range(100):
             node = random.randint(0, ts.num_nodes - 1)
-            pos = random.randint(0, ts.sequence_length - 1)
+            pos = random.randint(0, int(ts.sequence_length - 1))
             tree = ts.at(pos)
             parent = tree.parent(node)
             a = pyslim.mutation_at(ts, node, pos)
@@ -664,7 +664,7 @@ def test_nucleotide_at(self, recipe):
                 assert len(ts.reference_sequence.data) == ts.sequence_length
                 for _ in range(100):
                     node = random.randint(0, ts.num_nodes - 1)
-                    pos = random.randint(0, ts.sequence_length - 1)
+                    pos = random.randint(0, int(ts.sequence_length - 1))
                     tree = ts.at(pos)
                     parent = tree.parent(node)
                     a = pyslim.nucleotide_at(ts, node, pos)