add proper fst initialization, add copying of the fst

Author: bygu4

pyformlang/fst/__init__.py

@@ -12,10 +12,11 @@
 
 """
 
-from .fst import FST, State, Symbol, Epsilon
+from .fst import FST, TransitionFunction, State, Symbol, Epsilon
 
 
 __all__ = ["FST",
+           "TransitionFunction",
            "State",
            "Symbol",
            "Epsilon"]

pyformlang/fst/fst.py

@@ -1,36 +1,38 @@
 """ Finite State Transducer """
 
-from typing import Dict, List, Set, Tuple, Iterator, Iterable, Hashable
-from copy import deepcopy
+from typing import Dict, List, Set, AbstractSet, \
+    Tuple, Iterator, Iterable, Hashable
 
 from networkx import MultiDiGraph
 from networkx.drawing.nx_pydot import write_dot
 
+from .transition_function import TransitionFunction
+from .transition_function import TransitionKey, TransitionValues, Transition
 from .utils import StateRenaming
 from ..objects.finite_automaton_objects import State, Symbol, Epsilon
 from ..objects.finite_automaton_objects.utils import to_state, to_symbol
 
-TransitionKey = Tuple[State, Symbol]
-TransitionValue = Tuple[State, Tuple[Symbol, ...]]
-TransitionValues = Set[TransitionValue]
-TransitionFunction = Dict[TransitionKey, TransitionValues]
-
 InputTransition = Tuple[Hashable, Hashable, Hashable, Iterable[Hashable]]
-Transition = Tuple[TransitionKey, TransitionValue]
 
 
 class FST(Iterable[Transition]):
     """ Representation of a Finite State Transducer"""
 
-    def __init__(self) -> None:
-        self._states: Set[State] = set()  # Set of states
-        self._input_symbols: Set[Symbol] = set()  # Set of input symbols
-        self._output_symbols: Set[Symbol] = set()  # Set of output symbols
-        # Dict from _states x _input_symbols U {epsilon} into a subset of
-        # _states X _output_symbols*
-        self._delta: TransitionFunction = {}
-        self._start_states: Set[State] = set()
-        self._final_states: Set[State] = set()  # _final_states is final states
+    def __init__(self,
+                 states: AbstractSet[Hashable] = None,
+                 input_symbols: AbstractSet[Hashable] = None,
+                 output_symbols: AbstractSet[Hashable] = None,
+                 transition_function: TransitionFunction = None,
+                 start_states: AbstractSet[Hashable] = None,
+                 final_states: AbstractSet[Hashable] = None) -> None:
+        self._states = {to_state(x) for x in states or set()}
+        self._input_symbols = {to_symbol(x) for x in input_symbols or set()}
+        self._output_symbols = {to_symbol(x) for x in output_symbols or set()}
+        self._transition_function = transition_function or TransitionFunction()
+        self._start_states = {to_state(x) for x in start_states or set()}
+        self._states.update(self._start_states)
+        self._final_states = {to_state(x) for x in final_states or set()}
+        self._states.update(self._final_states)
 
     @property
     def states(self) -> Set[State]:
@@ -87,16 +89,6 @@ def final_states(self) -> Set[State]:
         """
         return self._final_states
 
-    def get_number_transitions(self) -> int:
-        """ Get the number of transitions in the FST
-
-        Returns
-        ----------
-        n_transitions : int
-            The number of transitions
-        """
-        return sum(len(x) for x in self._delta.values())
-
     def add_transition(self,
                        s_from: Hashable,
                        input_symbol: Hashable,
@@ -125,11 +117,10 @@ def add_transition(self,
         if input_symbol != Epsilon():
             self._input_symbols.add(input_symbol)
         self._output_symbols.update(output_symbols)
-        head = (s_from, input_symbol)
-        if head in self._delta:
-            self._delta[head].add((s_to, output_symbols))
-        else:
-            self._delta[head] = {(s_to, output_symbols)}
+        self._transition_function.add_transition(s_from,
+                                                 input_symbol,
+                                                 s_to,
+                                                 output_symbols)
 
     def add_transitions(self, transitions: Iterable[InputTransition]) -> None:
         """
@@ -156,8 +147,20 @@ def remove_transition(self,
         input_symbol = to_symbol(input_symbol)
         s_to = to_state(s_to)
         output_symbols = tuple(to_symbol(x) for x in output_symbols)
-        head = (s_from, input_symbol)
-        self._delta.get(head, set()).discard((s_to, output_symbols))
+        self._transition_function.remove_transition(s_from,
+                                                    input_symbol,
+                                                    s_to,
+                                                    output_symbols)
+
+    def get_number_transitions(self) -> int:
+        """ Get the number of transitions in the FST
+
+        Returns
+        ----------
+        n_transitions : int
+            The number of transitions
+        """
+        return self._transition_function.get_number_transitions()
 
     def add_start_state(self, start_state: Hashable) -> None:
         """ Add a start state
@@ -188,7 +191,7 @@ def __call__(self, s_from: Hashable, input_symbol: Hashable) \
         """ Calls the transition function of the FST """
         s_from = to_state(s_from)
         input_symbol = to_symbol(input_symbol)
-        return self._delta.get((s_from, input_symbol), set())
+        return self._transition_function(s_from, input_symbol)
 
     def __contains__(self, transition: InputTransition) -> bool:
         """ Whether the given transition is present in the FST """
@@ -201,9 +204,7 @@ def __contains__(self, transition: InputTransition) -> bool:
 
     def __iter__(self) -> Iterator[Transition]:
         """ Gets an iterator of transitions of the FST """
-        for key, values in self._delta.items():
-            for value in values:
-                yield key, value
+        yield from self._transition_function
 
     def translate(self,
                   input_word: Iterable[Hashable],
@@ -300,14 +301,12 @@ def _add_transitions_to(self,
                             union_fst: "FST",
                             state_renaming: StateRenaming,
                             idx: int) -> None:
-        for head, transition in self._delta.items():
-            s_from, input_symbol = head
-            for s_to, output_symbols in transition:
-                union_fst.add_transition(
-                    state_renaming.get_renamed_state(s_from, idx),
-                    input_symbol,
-                    state_renaming.get_renamed_state(s_to, idx),
-                    output_symbols)
+        for (s_from, input_symbol), (s_to, output_symbols) in self:
+            union_fst.add_transition(
+                state_renaming.get_renamed_state(s_from, idx),
+                input_symbol,
+                state_renaming.get_renamed_state(s_to, idx),
+                output_symbols)
 
     def _add_extremity_states_to(self,
                                  union_fst: "FST",
@@ -502,6 +501,18 @@ def write_as_dot(self, filename: str) -> None:
         """
         write_dot(self.to_networkx(), filename)
 
-    def to_dict(self) -> TransitionFunction:
+    def copy(self) -> "FST":
+        """ Copies the FST """
+        return FST(states=self.states,
+                   input_symbols=self.input_symbols,
+                   output_symbols=self.output_symbols,
+                   transition_function=self._transition_function.copy(),
+                   start_states=self.start_states,
+                   final_states=self.final_states)
+
+    def __copy__(self) -> "FST":
+        return self.copy()
+
+    def to_dict(self) -> Dict[TransitionKey, TransitionValues]:
         """Gives the transitions as a dictionary"""
-        return deepcopy(self._delta)
+        return self._transition_function.to_dict()

pyformlang/fst/tests/test_fst.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from pyformlang.fst import FST
+from pyformlang.fst import FST, TransitionFunction, State, Symbol
 
 
 @pytest.fixture
@@ -194,9 +194,9 @@ def test_paper(self):
     def test_contains(self, fst0: FST):
         """ Tests the containment of transition in the FST """
         assert ("q0", "a", "q1", ["b"]) in fst0
-        assert ("a", "b", "c", "d") not in fst0
-        fst0.add_transition("a", "b", "c", "d")
-        assert ("a", "b", "c", "d") in fst0
+        assert ("a", "b", "c", ["d"]) not in fst0
+        fst0.add_transition("a", "b", "c", {"d"})
+        assert ("a", "b", "c", ["d"]) in fst0
 
     def test_iter(self, fst0: FST):
         """ Tests the iteration of FST transitions """
@@ -210,9 +210,42 @@ def test_iter(self, fst0: FST):
 
     def test_remove_transition(self, fst0: FST):
         """ Tests the removal of transition from the FST """
-        assert ("q0", "a", "q1", "b") in fst0
-        fst0.remove_transition("q0", "a", "q1", "b")
-        assert ("q0", "a", "q1", "b") not in fst0
-        fst0.remove_transition("q0", "a", "q1", "b")
-        assert ("q0", "a", "q1", "b") not in fst0
+        assert ("q0", "a", "q1", ["b"]) in fst0
+        fst0.remove_transition("q0", "a", "q1", ["b"])
+        assert ("q0", "a", "q1", ["b"]) not in fst0
+        fst0.remove_transition("q0", "a", "q1", ["b"])
+        assert ("q0", "a", "q1", ["b"]) not in fst0
         assert fst0.get_number_transitions() == 0
+
+    def test_initialization(self):
+        """ Tests the initialization of the FST """
+        fst = FST(states={0},
+                  input_symbols={"a", "b"},
+                  output_symbols={"c"},
+                  start_states={1},
+                  final_states={2})
+        assert fst.states == {0, 1, 2}
+        assert fst.input_symbols == {"a", "b"}
+        assert fst.output_symbols == {"c"}
+        assert fst.get_number_transitions() == 0
+        assert not list(iter(fst))
+
+        function = TransitionFunction()
+        function.add_transition(State(1), Symbol("a"), State(2), (Symbol("b"),))
+        function.add_transition(State(1), Symbol("a"), State(2), (Symbol("c"),))
+        fst = FST(transition_function=function)
+        assert fst.get_number_transitions() == 2
+        assert (1, "a", 2, ["b"]) in fst
+        assert (1, "a", 2, ["c"]) in fst
+        assert fst(1, "a") == {(2, tuple("b")), (2, tuple("c"))}
+
+    def test_copy(self, fst0: FST):
+        """ Tests the copying of the FST """
+        fst_copy = fst0.copy()
+        assert fst_copy.states == fst0.states
+        assert fst_copy.input_symbols == fst0.input_symbols
+        assert fst_copy.output_symbols == fst0.output_symbols
+        assert fst_copy.start_states == fst0.start_states
+        assert fst_copy.final_states == fst0.final_states
+        assert fst_copy.to_dict() == fst0.to_dict()
+        assert fst_copy is not fst0

pyformlang/fst/transition_function.py

@@ -0,0 +1,78 @@
+""" The transition function of Finite State Transducer """
+
+from typing import Dict, Set, Tuple, Iterator, Iterable
+from copy import deepcopy
+
+from ..objects.finite_automaton_objects import State, Symbol
+
+TransitionKey = Tuple[State, Symbol]
+TransitionValue = Tuple[State, Tuple[Symbol, ...]]
+TransitionValues = Set[TransitionValue]
+Transition = Tuple[TransitionKey, TransitionValue]
+
+
+class TransitionFunction(Iterable[Transition]):
+    """ The transition function of Finite State Transducer """
+
+    def __init__(self) -> None:
+        self._transitions: Dict[TransitionKey, TransitionValues] = {}
+
+    def add_transition(self,
+                       s_from: State,
+                       input_symbol: Symbol,
+                       s_to: State,
+                       output_symbols: Tuple[Symbol, ...]) -> None:
+        """ Adds given transition to the function """
+        key = (s_from, input_symbol)
+        value = (s_to, output_symbols)
+        self._transitions.setdefault(key, set()).add(value)
+
+    def remove_transition(self,
+                          s_from: State,
+                          input_symbol: Symbol,
+                          s_to: State,
+                          output_symbols: Tuple[Symbol, ...]) -> None:
+        """ Removes given transition from the function """
+        key = (s_from, input_symbol)
+        value = (s_to, output_symbols)
+        self._transitions.get(key, set()).discard(value)
+
+    def get_number_transitions(self) -> int:
+        """ Gets the number of transitions in the function
+
+        Returns
+        ----------
+        n_transitions : int
+            The number of transitions
+        """
+        return sum(len(x) for x in self._transitions.values())
+
+    def __call__(self, s_from: State, input_symbol: Symbol) \
+            -> TransitionValues:
+        """ Calls the transition function """
+        return self._transitions.get((s_from, input_symbol), set())
+
+    def __contains__(self, transition: Transition) -> bool:
+        """ Whether the given transition is present in the function """
+        key, value = transition
+        return value in self(*key)
+
+    def __iter__(self) -> Iterator[Transition]:
+        """ Gets an iterator of transitions of the function """
+        for key, values in self._transitions.items():
+            for value in values:
+                yield key, value
+
+    def copy(self) -> "TransitionFunction":
+        """ Copies the transition function """
+        new_tf = TransitionFunction()
+        for key, value in self:
+            new_tf.add_transition(*key, *value)
+        return new_tf
+
+    def __copy__(self) -> "TransitionFunction":
+        return self.copy()
+
+    def to_dict(self) -> Dict[TransitionKey, TransitionValues]:
+        """ Gives the transition function as a dictionary """
+        return deepcopy(self._transitions)