Check tokens configuration

dan/utils.py

@@ -2,6 +2,7 @@
 import json
 from argparse import ArgumentTypeError
 from itertools import islice
+from operator import attrgetter
 from pathlib import Path
 from typing import Dict, NamedTuple
 
@@ -131,10 +132,20 @@ def list_to_batches(iterable, n):
 
 
 def parse_tokens(filename: str) -> Dict[str, EntityType]:
-    return {
+    tokens = {
         name: EntityType(**tokens)
         for name, tokens in yaml.safe_load(Path(filename).read_text()).items()
     }
+    end_tokens = map(attrgetter("end"), tokens.values())
+    # Check that either
+    if next(end_tokens):
+        # - all entities have end token
+        assert all(end_tokens), "Some entities have no end token"
+    else:
+        # - no entities have end tokens
+        assert not any(end_tokens), "Some entities have end tokens"
+
+    return tokens
 
 
 def read_yaml(yaml_path: str) -> Dict:

tests/test_utils.py

+import pytest
+import yaml
+
+from dan.utils import parse_tokens
+
+
+@pytest.mark.parametrize(
+    ("tokens", "error_msg"),
+    [
+        # All should have no end tokens
+        (
+            {
+                "name": {
+                    "start": "A",
+                    "end": "",
+                },
+                "surname": {"start": "B", "end": "C"},
+            },
+            "Some entities have end tokens",
+        ),
+        # All should have end tokens
+        (
+            {
+                "name": {
+                    "start": "A",
+                    "end": "C",
+                },
+                "surname": {"start": "B", "end": ""},
+            },
+            "Some entities have no end token",
+        ),
+    ],
+)
+def test_parse_tokens_errors(tmp_path, tokens, error_msg):
+    tokens_path = tmp_path / "tokens.yml"
+    tokens_path.write_text(yaml.dump(tokens))
+
+    with pytest.raises(AssertionError, match=error_msg):
+        parse_tokens(tokens_path)