Reformat add_docs_equal and add docstring

This commit is contained in:
Ines Montani 2017-01-13 14:25:53 +01:00
parent 62ce0a0073
commit db9b25663d
1 changed files with 2 additions and 4 deletions

View File

@ -46,20 +46,18 @@ def get_cosine(vec1, vec2):
def assert_docs_equal(doc1, doc2):
# tokens
"""Compare two Doc objects and assert that they're equal. Tests for tokens,
tags, dependencies and entities."""
assert [ t.orth for t in doc1 ] == [ t.orth for t in doc2 ]
# tags
assert [ t.pos for t in doc1 ] == [ t.pos for t in doc2 ]
assert [ t.tag for t in doc1 ] == [ t.tag for t in doc2 ]
# parse
assert [ t.head.i for t in doc1 ] == [ t.head.i for t in doc2 ]
assert [ t.dep for t in doc1 ] == [ t.dep for t in doc2 ]
if doc1.is_parsed and doc2.is_parsed:
assert [ s for s in doc1.sents ] == [ s for s in doc2.sents ]
# entities
assert [ t.ent_type for t in doc1 ] == [ t.ent_type for t in doc2 ]
assert [ t.ent_iob for t in doc1 ] == [ t.ent_iob for t in doc2 ]
assert [ ent for ent in doc1.ents ] == [ ent for ent in doc2.ents ]