spaCy/spacy/tests/pipeline/test_factories.py

# coding: utf8
from __future__ import unicode_literals

import pytest
from spacy.language import Language
from spacy.tokens import Span

from ..util import get_doc


@pytest.fixture
def doc(en_tokenizer):
    text = 'I like New York in Autumn.'
    heads = [1, 0, 1, -2, -3, -1, -5]
    tags = ['PRP', 'IN', 'NNP', 'NNP', 'IN', 'NNP', '.']
    pos = ['PRON', 'VERB', 'PROPN', 'PROPN', 'ADP', 'PROPN', 'PUNCT']
    deps = ['ROOT', 'prep', 'compound', 'pobj', 'prep', 'pobj', 'punct']
    tokens = en_tokenizer(text)
    doc = get_doc(tokens.vocab, words=[t.text for t in tokens], heads=heads,
                  tags=tags, pos=pos, deps=deps)
    doc.ents = [Span(doc, 2, 4, doc.vocab.strings['GPE'])]
    doc.is_parsed = True
    doc.is_tagged = True
    return doc


def test_factories_merge_noun_chunks(doc):
    assert len(doc) == 7
    nlp = Language()
    merge_noun_chunks = nlp.create_pipe('merge_noun_chunks')
    merge_noun_chunks(doc)
    assert len(doc) == 6
    assert doc[2].text == 'New York'


def test_factories_merge_ents(doc):
    assert len(doc) == 7
    assert len(list(doc.ents)) == 1
    nlp = Language()
    merge_entities = nlp.create_pipe('merge_entities')
    merge_entities(doc)
    assert len(doc) == 6
    assert len(list(doc.ents)) == 1
    assert doc[2].text == 'New York'