def test_exclude_types(self, spacy_doc): ent_types = ["PERSON", "GPE"] for exclude_types in ent_types: result = extract.entities(spacy_doc, exclude_types=exclude_types) assert all(span.label_ != exclude_types for span in result) ent_types = [{"PERSON", "GPE"}, ("DATE", "ORG"), ["LOC"]] for exclude_types in ent_types: result = extract.entities(spacy_doc, exclude_types=exclude_types) assert all(span.label_ not in exclude_types for span in result) # special numeric cases! ent_types = ["NUMERIC", ("NUMERIC",), {"PERSON", "NUMERIC"}] for exclude_types in ent_types: exclude_types_parsed = extract._parse_ent_types(exclude_types, "exclude") result = extract.entities(spacy_doc, exclude_types=exclude_types) assert all(span.label_ not in exclude_types_parsed for span in result)
def test_parse_ent_types_bad_type(self): for bad_type in [1, 3.1415, True, b"PERSON"]: with pytest.raises(TypeError): _ = extract._parse_ent_types(bad_type, "include")