Successors length limits (#224)

* Succesor length limits

* Added extra test cases and comment about bug

* Fix infinite loop bug when length is limited.

* Re-enable disabled test

* Add an edge case test for predecessor

* Formatting

---------

Co-authored-by: Eliot Robson <eliot.robson24@gmail.com>

Author: Tagl

automata/fa/dfa.py

@@ -1276,6 +1276,8 @@ def predecessor(
         *,
         strict: bool = True,
         key: Optional[Callable[[Any], Any]] = None,
+        min_length: int = 0,
+        max_length: Optional[int] = None,
     ) -> Optional[str]:
         """
         Returns the first string accepted by the DFA that comes before
@@ -1291,6 +1293,10 @@ def predecessor(
         key : Optional[Callable], default: None
             Function for defining custom lexicographical ordering. Defaults to using
             the standard string ordering.
+        min_length : int, default: 0
+            Limits generation to words with at least the given length.
+        max_length : Optional[int], default: None
+            Limits generation to words with at most the given length.
 
         Returns
         ------
@@ -1303,7 +1309,13 @@ def predecessor(
             Raised if the language accepted by self is infinite, as we cannot
             generate predecessors in this case.
         """
-        for word in self.predecessors(input_str, strict=strict, key=key):
+        for word in self.predecessors(
+            input_str,
+            strict=strict,
+            key=key,
+            min_length=min_length,
+            max_length=max_length,
+        ):
             return word
         return None
 
@@ -1313,6 +1325,8 @@ def predecessors(
         *,
         strict: bool = True,
         key: Optional[Callable[[Any], Any]] = None,
+        min_length: int = 0,
+        max_length: Optional[int] = None,
     ) -> Generator[str, None, None]:
         """
         Generates all strings that come before the input string
@@ -1328,6 +1342,10 @@ def predecessors(
         key : Optional[Callable], default: None
             Function for defining custom lexicographical ordering. Defaults to using
             the standard string ordering.
+        min_length : int, default: 0
+            Limits generation to words with at least the given length.
+        max_length : Optional[int], default: None
+            Limits generation to words with at most the given length.
 
         Returns
         ------
@@ -1341,14 +1359,23 @@ def predecessors(
             Raised if the language accepted by self is infinite, as we cannot
             generate predecessors in this case.
         """
-        yield from self.successors(input_str, strict=strict, reverse=True, key=key)
+        yield from self.successors(
+            input_str,
+            strict=strict,
+            reverse=True,
+            key=key,
+            min_length=min_length,
+            max_length=max_length,
+        )
 
     def successor(
         self,
         input_str: Optional[str],
         *,
         strict: bool = True,
         key: Optional[Callable[[Any], Any]] = None,
+        min_length: int = 0,
+        max_length: Optional[int] = None,
     ) -> Optional[str]:
         """
         Returns the first string accepted by the DFA that comes after
@@ -1364,13 +1391,23 @@ def successor(
         key : Optional[Callable], default: None
             Function for defining custom lexicographical ordering. Defaults to using
             the standard string ordering.
+        min_length : int, default: 0
+            Limits generation to words with at least the given length.
+        max_length : Optional[int], default: None
+            Limits generation to words with at most the given length.
 
         Returns
         ------
         str
             The first string accepted by the DFA lexicographically before input_string.
         """
-        for word in self.successors(input_str, strict=strict, key=key):
+        for word in self.successors(
+            input_str,
+            strict=strict,
+            key=key,
+            min_length=min_length,
+            max_length=max_length,
+        ):
             return word
         return None
 
@@ -1381,6 +1418,8 @@ def successors(
         strict: bool = True,
         key: Optional[Callable[[Any], Any]] = None,
         reverse: bool = False,
+        min_length: int = 0,
+        max_length: Optional[int] = None,
     ) -> Generator[str, None, None]:
         """
         Generates all strings that come after the input string
@@ -1398,6 +1437,10 @@ def successors(
             the standard string ordering.
         reverse : bool, default: False
             If True, then predecessors will be generated instead of successors.
+        min_length : int, default: 0
+            Limits generation to words with at least the given length.
+        max_length : Optional[int], default: None
+            Limits generation to words with at most the given length.
 
         Returns
         ------
@@ -1446,6 +1489,8 @@ def successors(
             if (
                 not reverse
                 and should_yield
+                and min_length <= len(char_stack)
+                and (max_length is None or len(char_stack) <= max_length)
                 and candidate == first_symbol
                 and state in self.final_states
             ):
@@ -1456,7 +1501,9 @@ def successors(
                 else self._get_next_current_state(state, candidate)
             )
             # Traverse to child if candidate is viable
-            if candidate_state in coaccessible_nodes:
+            if candidate_state in coaccessible_nodes and (
+                max_length is None or len(char_stack) < max_length
+            ):
                 state_stack.append(candidate_state)
                 char_stack.append(cast(str, candidate))
                 candidate = first_symbol
@@ -1465,6 +1512,8 @@ def successors(
                 if (
                     reverse
                     and should_yield
+                    and min_length <= len(char_stack)
+                    and (max_length is None or len(char_stack) <= max_length)
                     and candidate is None
                     and state in self.final_states
                 ):
@@ -1480,6 +1529,8 @@ def successors(
         if (
             reverse
             and should_yield
+            and min_length <= len(char_stack)
+            and (max_length is None or len(char_stack) <= max_length)
             and candidate is None
             and state in self.final_states
         ):

tests/test_dfa.py

@@ -1836,6 +1836,9 @@ def test_predecessor(self, as_partial: bool) -> None:
         actual = list(dfa.predecessors("010", strict=False))
 
         self.assertEqual(dfa.predecessor("000"), "00")
+        self.assertEqual(dfa.predecessor("000", max_length=1), "0")
+        self.assertEqual(dfa.predecessor("0", min_length=2), None)
+        self.assertEqual(dfa.predecessor("0000", min_length=2, max_length=3), "000")
         self.assertEqual(dfa.predecessor("0100"), "010")
         self.assertEqual(dfa.predecessor("1"), "010101111111101011010100")
         self.assertEqual(
@@ -1872,6 +1875,10 @@ def test_successor(self, as_partial: bool) -> None:
         self.assertIsNone(dfa.successor("110"))
         self.assertIsNone(dfa.successor("111111110101011"))
 
+        self.assertEqual(dfa.successor("", min_length=3), "000")
+        self.assertEqual(dfa.successor("", min_length=4), "010101111111101011010100")
+        self.assertEqual(dfa.successor("010", max_length=6), "100")
+
         infinite_dfa = DFA.from_nfa(NFA.from_regex("0*1*"))
         self.assertEqual(infinite_dfa.successor(""), "0")
         self.assertEqual(infinite_dfa.successor("0"), "00")
@@ -1882,6 +1889,10 @@ def test_successor(self, as_partial: bool) -> None:
         self.assertEqual(infinite_dfa.successor("1"), "11")
         self.assertEqual(infinite_dfa.successor(100 * "0"), 101 * "0")
         self.assertEqual(infinite_dfa.successor(100 * "1"), 101 * "1")
+        self.assertEqual(infinite_dfa.successor("", min_length=5), "00000")
+        self.assertEqual(infinite_dfa.successor("000", min_length=5), "00000")
+        self.assertEqual(infinite_dfa.successor("1", min_length=5), "11111")
+        self.assertEqual(infinite_dfa.successor("1111", max_length=4), None)
 
     @params(True, False)
     def test_successor_and_predecessor(self, as_partial: bool) -> None: