fsm.py

from collections import defaultdict


class DFA:
    def __init__(self, initial_state, final_state, rules):
        self.initial_state = initial_state
        self.final_state = final_state
        self.rules = rules

        self.transitions = defaultdict(dict)
        self.alphabet = set()
        self.states = set()
        for current, input, next in rules:
            self.transitions[current][input] = next
            self.alphabet.add(input)
            self.states.add(current)
            self.states.add(next)

    def next_state(self, state, input):
        return self.transitions[state][input]

    def accept(self, s):
        state = self.initial_state
        for c in s:
            state = self.next_state(state, c)
        return state in self.final_state


class NFA:
    def __init__(self, initial_state, final_state, rules):
        self.initial_state = initial_state
        self.final_state = final_state
        self.rules = rules

        self.transitions = defaultdict(lambda: defaultdict(set))
        self.alphabet = set()
        self.states = set()
        for current, input, next in rules:
            self.transitions[current][input].add(next)
            self.alphabet.add(input)
            self.states.add(current)
            self.states.add(next)

    def next_states(self, state, input):
        return self.transitions[state][input]

    def accept(self, s):
        current_states = {self.initial_state}
        for c in s:
            next_states = set()
            for state in current_states:
                next_states |= self.next_states(state, c)
            current_states = next_states
        return not current_states.isdisjoint(self.final_state)

test_fsm.py

import pytest

class TestDFA:
    @pytest.fixture
    def target(self, *args, **kwargs):
        import fsm

        return fsm.DFA

    @pytest.mark.parametrize("state,input,expected", [
        (1, 'a', 2),
        (2, 'a', 2),
        (2, 'a', 2),
        (3, 'a', 3),
    ])
    def test_next_state(self, target, state, input, expected):
        dfa = target(1, {3}, [
            (1, 'a', 2),
            (1, 'b', 1),
            (2, 'a', 2),
            (2, 'b', 3),
            (3, 'a', 3),
            (3, 'b', 3),
        ])
        assert dfa.next_state(state, input) == expected

    @pytest.mark.parametrize("input,expected", [
        ('a', False),
        ('aa', False),
        ('baa', False),
        ('baba', True),
    ])
    def test_accept(self, target, input, expected):
        dfa = target(1, {3}, [
            (1, 'a', 2),
            (1, 'b', 1),
            (2, 'a', 2),
            (2, 'b', 3),
            (3, 'a', 3),
            (3, 'b', 3),
        ])
        assert dfa.accept(input) == expected


class TestNFA:
    @pytest.fixture
    def target(self, *args, **kwargs):
        import fsm

        return fsm.NFA

    @pytest.mark.parametrize("state,input,expected", [
        (1, 'a', {1}),
        (1, 'b', {1, 2}),        
        (2, 'a', {3}),
        (2, 'b', {3}),        
    ])
    def test_next_states(self, target, state, input, expected):
        nfa = target(1, {3}, [
            (1, 'a', 1),
            (1, 'b', 1),
            (1, 'b', 2),
            (2, 'a', 3),
            (2, 'b', 3),
        ])
        assert nfa.next_states(state, input) == expected

    @pytest.mark.parametrize("input,expected", [
        ('a', False),
        ('aa', False),
        ('ab', False),
        ('b', False),
        ('abba', True),
        ('aababbb', True)
    ])
    def test_accept(self, target, input, expected):
        nfa = target(1, {3}, [
            (1, 'a', 1),
            (1, 'b', 1),
            (1, 'b', 2),
            (2, 'a', 3),
            (2, 'b', 3),
        ])
        assert nfa.accept(input) == expected