Fix pickle tests

This commit is contained in:
Matthew Honnibal 2018-12-06 20:43:47 +01:00
parent 0a60726215
commit e619f45287

View File

@ -2,8 +2,8 @@
from __future__ import unicode_literals from __future__ import unicode_literals
import pytest import pytest
import dill as pickle
import numpy import numpy
import srsly
from spacy.strings import StringStore from spacy.strings import StringStore
from spacy.vocab import Vocab from spacy.vocab import Vocab
from spacy.attrs import NORM from spacy.attrs import NORM
@ -14,8 +14,8 @@ def test_pickle_string_store(text1, text2):
stringstore = StringStore() stringstore = StringStore()
store1 = stringstore[text1] store1 = stringstore[text1]
store2 = stringstore[text2] store2 = stringstore[text2]
data = pickle.dumps(stringstore, protocol=-1) data = srsly.pickle_dumps(stringstore, protocol=-1)
unpickled = pickle.loads(data) unpickled = srsly.pickle_loads(data)
assert unpickled[text1] == store1 assert unpickled[text1] == store1
assert unpickled[text2] == store2 assert unpickled[text2] == store2
assert len(stringstore) == len(unpickled) assert len(stringstore) == len(unpickled)