diff --git a/automata/base/automaton.py b/automata/base/automaton.py index af4cb38c..6580475f 100644 --- a/automata/base/automaton.py +++ b/automata/base/automaton.py @@ -39,6 +39,16 @@ def __delattr__(self, name): """Set custom delattr to make class immutable.""" raise AttributeError(f'This {type(self).__name__} is immutable') + def __getstate__(self): + """Return the object's state, described by its input parameters""" + return self.input_parameters + + def __setstate__(self, d): + """Restore the object state from its input parameters""" + # Notice that the default __setstate__ method won't work + # because __setattr__ is disabled due to immutability + self.__init__(**d) + @abc.abstractmethod def read_input_stepwise(self, input_str): """Return a generator that yields each step while reading input.""" diff --git a/tests/test_serialization.py b/tests/test_serialization.py new file mode 100644 index 00000000..624af174 --- /dev/null +++ b/tests/test_serialization.py @@ -0,0 +1,25 @@ +#!/usr/bin/env python3 +"""Tests for automata pickle serialization""" + +import pickle + +import tests.test_fa as test_fa + + +class TestSerialization(test_fa.TestFA): + """ + This tests verifies that the cycle FA -> .pkl -> FA works. + The test only applies to DFA and NFAs as other classes do not implement equality + """ + + def test_serialize_dfa(self): + """Should convert a DFA to pickle serialization and reads it back""" + s = pickle.dumps(self.dfa) + dfa = pickle.loads(s) + self.assertEqual(self.dfa, dfa) + + def test_serialize_nfa(self): + """Should convert a NFA to pickled representation and read it back""" + s = pickle.dumps(self.nfa) + nfa = pickle.loads(s) + self.assertEqual(self.nfa, nfa)