spaCy/spacy/tests/serialize/test_serialize_extension_attrs.py
Daniël de Kok 68d7841df5
Extension serialization attr tests: add teardown (#13284)
The doc/token extension serialization tests add extensions that are not
serializable with pickle. This didn't cause issues before due to the
implicit run order of tests. However, test ordering has changed with
pytest 8.0.0, leading to failed tests in test_language.

Update the fixtures in the extension serialization tests to do proper
teardown and remove the extensions.
2024-01-29 13:51:56 +01:00

36 lines
1.1 KiB
Python

import pytest
from spacy.tokens import Doc, Token
from spacy.vocab import Vocab
@pytest.fixture
def doc_w_attrs(en_tokenizer):
Doc.set_extension("_test_attr", default=False)
Doc.set_extension("_test_prop", getter=lambda doc: len(doc.text))
Doc.set_extension("_test_method", method=lambda doc, arg: f"{len(doc.text)}{arg}")
doc = en_tokenizer("This is a test.")
doc._._test_attr = "test"
Token.set_extension("_test_token", default="t0")
doc[1]._._test_token = "t1"
yield doc
Doc.remove_extension("_test_attr")
Doc.remove_extension("_test_prop")
Doc.remove_extension("_test_method")
Token.remove_extension("_test_token")
def test_serialize_ext_attrs_from_bytes(doc_w_attrs):
doc_b = doc_w_attrs.to_bytes()
doc = Doc(Vocab()).from_bytes(doc_b)
assert doc._.has("_test_attr")
assert doc._._test_attr == "test"
assert doc._._test_prop == len(doc.text)
assert doc._._test_method("test") == f"{len(doc.text)}test"
assert doc[0]._._test_token == "t0"
assert doc[1]._._test_token == "t1"
assert doc[2]._._test_token == "t0"