From 06a9939eb5f39614859ce1d201a85b11590f0726 Mon Sep 17 00:00:00 2001
From: Richard Hudson <richard@explosion.ai>
Date: Tue, 30 Nov 2021 17:06:55 +0100
Subject: [PATCH] Add dependency tree tests

---
 spacy/tests/test_visualization.py | 161 ++++++++++++++++++++++++++++++
 spacy/visualization.py            |   4 +
 2 files changed, 165 insertions(+)
 create mode 100644 spacy/tests/test_visualization.py

diff --git a/spacy/tests/test_visualization.py b/spacy/tests/test_visualization.py
new file mode 100644
index 000000000..7f64de54b
--- /dev/null
+++ b/spacy/tests/test_visualization.py
@@ -0,0 +1,161 @@
+import pytest
+from spacy.visualization import Visualizer
+from spacy.tokens import Span, Doc
+
+
+def test_dependency_tree_basic(en_vocab):
+    """Test basic dependency tree display."""
+    doc = Doc(
+        en_vocab,
+        words=[
+            "The",
+            "big",
+            "dog",
+            "chased",
+            "the",
+            "frightened",
+            "cat",
+            "mercilessly",
+            ".",
+        ],
+        heads=[2, 2, 3, None, 6, 6, 3, 3, 3],
+        deps=["dep"] * 9,
+    )
+    dep_tree = Visualizer.render_dependency_tree(doc[0 : len(doc)], True)
+    assert dep_tree == [
+        "<╗  ",
+        "<╣  ",
+        "═╝<╗",
+        "═══╣",
+        "<╗ ║",
+        "<╣ ║",
+        "═╝<╣",
+        "<══╣",
+        "<══╝",
+    ]
+    dep_tree = Visualizer.render_dependency_tree(doc[0 : len(doc)], False)
+    assert dep_tree == [
+        "  ╔>",
+        "  ╠>",
+        "╔>╚═",
+        "╠═══",
+        "║ ╔>",
+        "║ ╠>",
+        "╠>╚═",
+        "╠══>",
+        "╚══>",
+    ]
+
+
+def test_dependency_tree_non_initial_sentence(en_vocab):
+    """Test basic dependency tree display."""
+    doc = Doc(
+        en_vocab,
+        words=[
+            "Something",
+            "happened",
+            ".",
+            "The",
+            "big",
+            "dog",
+            "chased",
+            "the",
+            "frightened",
+            "cat",
+            "mercilessly",
+            ".",
+        ],
+        heads=[0, None, 0, 5, 5, 6, None, 9, 9, 6, 6, 6],
+        deps=["dep"] * 12,
+    )
+    dep_tree = Visualizer.render_dependency_tree(doc[3 : len(doc)], True)
+    assert dep_tree == [
+        "<╗  ",
+        "<╣  ",
+        "═╝<╗",
+        "═══╣",
+        "<╗ ║",
+        "<╣ ║",
+        "═╝<╣",
+        "<══╣",
+        "<══╝",
+    ]
+    dep_tree = Visualizer.render_dependency_tree(doc[3 : len(doc)], False)
+    assert dep_tree == [
+        "  ╔>",
+        "  ╠>",
+        "╔>╚═",
+        "╠═══",
+        "║ ╔>",
+        "║ ╠>",
+        "╠>╚═",
+        "╠══>",
+        "╚══>",
+    ]
+
+
+def test_dependency_tree_non_projective(en_vocab):
+    """Test dependency tree display with a non-prejective dependency."""
+    doc = Doc(
+        en_vocab,
+        words=[
+            "I",
+            "saw",
+            "a",
+            "horse",
+            "yesterday",
+            "that",
+            "was",
+            "injured",
+            ".",
+        ],
+        heads=[1, None, 3, 1, 1, 7, 7, 3, 1],
+        deps=["dep"] * 9,
+    )
+    dep_tree = Visualizer.render_dependency_tree(doc[0 : len(doc)], True)
+    assert dep_tree == [
+        "<╗    ",
+        "═╩═══╗",
+        "<╗   ║",
+        "═╩═╗<╣",
+        "<══║═╣",
+        "<╗ ║ ║",
+        "<╣ ║ ║",
+        "═╝<╝ ║",
+        "<════╝",
+    ]
+    dep_tree = Visualizer.render_dependency_tree(doc[0 : len(doc)], False)
+    assert dep_tree == [
+        "    ╔>",
+        "╔═══╩═",
+        "║   ╔>",
+        "╠>╔═╩═",
+        "╠═║══>",
+        "║ ║ ╔>",
+        "║ ║ ╠>",
+        "║ ╚>╚═",
+        "╚════>",
+    ]
+
+
+def test_dependency_tree_input_not_span(en_vocab):
+    """Test dependency tree display behaviour when the input is not a Span."""
+    doc = Doc(
+        en_vocab,
+        words=[
+            "I",
+            "saw",
+            "a",
+            "horse",
+            "yesterday",
+            "that",
+            "was",
+            "injured",
+            ".",
+        ],
+        heads=[1, None, 3, 1, 1, 7, 7, 3, 1],
+        deps=["dep"] * 9,
+    )
+
+    with pytest.raises(AssertionError):
+        Visualizer.render_dependency_tree(doc[1:3], True)
diff --git a/spacy/visualization.py b/spacy/visualization.py
index 5d9dd04a2..13cf3693b 100644
--- a/spacy/visualization.py
+++ b/spacy/visualization.py
@@ -50,6 +50,10 @@ class Visualizer:
 
         Adapted from https://github.com/KoichiYasuoka/deplacy
         """
+
+        # Check sent is really a sentence
+        assert sent.start == sent[0].sent.start
+        assert sent.end == sent[0].sent.end
         heads = [
             None
             if token.dep_.lower() == "root" or token.head.i == token.i