import copy
import pickle

from flair.data import Sentence

sentence: Sentence = Sentence("Berlin is the capital of Germany.")

# Works: pickle / copy on unannotated sentence
pickled: Sentence = pickle.loads(pickle.dumps(sentence))
copied: Sentence = copy.deepcopy(sentence)
assert pickled.to_original_text() == copied.to_original_text() == "Berlin is the capital of Germany."

# Add span annotations
sentence[:1].add_label("ner", value="LOC")  # Berlin
sentence[5:6].add_label("ner", value="LOC")  # Germany

# Does not work: pickle / copy on span-annotated sentence
# TypeError: __new__() missing 1 required positional argument: 'tokens'
pickled: Sentence = pickle.loads(pickle.dumps(sentence))
copied: Sentence = copy.deepcopy(sentence)
assert (
    {span.text for span in pickled.get_spans("ner")}
    == {span.text for span in copied.get_spans("ner")}
    == {"Berlin", "Germany"}
)