Optimize backreference handling

Author: Buba98

regex_enumerator/regex_enumerator.py

@@ -4,7 +4,12 @@
 
 class RegexEnumerator:
     def __init__(self, regex: str, additional_charset: str | list[str] = None) -> None:
-        default_charset = [chr(c) for c in range(32, 127)]
+        default_charset = [' ', '!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+', ',', '-', '.', '/',
+                           '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', ':', ';', '<', '=', '>', '?',
+                           '@', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O',
+                           'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '[', '\\', ']', '^', '_',
+                           '`', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o',
+                           'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '{', '|', '}', '~']
 
         if additional_charset is None:
             additional = []

regex_enumerator/regex_parser.py

@@ -82,17 +82,16 @@ def _parseRegex(self, to_close: bool) -> RegexTree:
                     elements.append(
                         CharClass(chars, min_len, max_len))
                 case '.':
-                    chars = list(self.charset)
                     min_len, max_len = self._parseQuantifier()
                     elements.append(
-                        CharClass(chars, min_len, max_len))
+                        CharClass(self.charset, min_len, max_len))
                 case '\\':
                     reference = self._parseBackReferenceLookahead()
                     if reference is None:
                         chars = self._parseEscapeChar()
                         min_len, max_len = self._parseQuantifier()
                         elements.append(
-                            CharClass([chars], min_len, max_len))
+                            CharClass(chars, min_len, max_len))
                         continue
                     if isinstance(reference, str):
                         if reference not in named_groups:
@@ -110,7 +109,7 @@ def _parseRegex(self, to_close: bool) -> RegexTree:
                 case _:
                     min_len, max_len = self._parseQuantifier()
                     elements.append(
-                        CharClass([char], min_len, max_len))
+                        CharClass(char, min_len, max_len))
 
         if to_close:
             self._raise_error("Unmatched opening parenthesis")
@@ -191,7 +190,7 @@ def _parseEscapeChar(self) -> str:
                 self._raise_error("Unicode property not supported")
             case _: return char
 
-    def _parseCharClass(self) -> list[str]:
+    def _parseCharClass(self) -> str:
         chars_list: list[str] = []
         first_char = None
         range_divider = False
@@ -249,11 +248,12 @@ def _parseCharClass(self) -> list[str]:
         elif first_char is not None:
             chars_list.append(first_char)
 
+        charset = ''.join(sorted(set(''.join(chars_list))))
+
         if negated:
-            chars_list = [
-                c for c in self.charset if c not in ''.join(chars_list)]
+            return ''.join(c for c in self.charset if c not in charset)
 
-        return chars_list
+        return charset
 
     def _parseQuantifier(self) -> tuple[int, int | None]:
 

regex_enumerator/regex_tree.py

@@ -3,12 +3,12 @@ class RegexTree:
 
 
 class CharClass:
-    def __init__(self, chars_list: list[str], min_len: int, max_len: int | None):
+    def __init__(self, charset: str, min_len: int, max_len: int | None):
         self._index = 0
-        self._chars: str = ''.join(sorted(set(''.join(chars_list))))
+        self._charset = charset
         self._min_len = min_len
         self._max_len = max_len
-        self._base = len(self._chars)
+        self._base = len(charset)
         self.done = self._base == 0 or self._max_len == 0
         self.current: list[str] = self._first()
 
@@ -18,7 +18,7 @@ def _first(self) -> list[str]:
 
         if self._base == 1 and self._max_len is not None:
             self.done = True
-            result = [self._chars *
+            result = [self._charset *
                       i for i in range(self._min_len, self._max_len + 1)]
             return result
 
@@ -27,7 +27,7 @@ def _first(self) -> list[str]:
 
         result = ['']
         for _ in range(self._min_len):
-            result = [pfx + sfx for pfx in self._chars for sfx in result]
+            result = [pfx + sfx for pfx in self._charset for sfx in result]
 
         self._last = result
         return result
@@ -39,7 +39,7 @@ def next(self) -> list[str]:
         if self._max_len is not None and self._index + self._min_len == self._max_len:
             self.done = True
 
-        result = [pfx + sfx for pfx in self._last for sfx in self._chars]
+        result = [pfx + sfx for pfx in self._last for sfx in self._charset]
         self.current.extend(result)
         self._last = result
         return result
@@ -118,26 +118,18 @@ def next(self) -> set[str]:
         self._index = index
         result: list[tuple[str, dict[RegexTree, str]]] = []
 
-        if isinstance(self._elements[0], CharClass):
+        if isinstance(self._elements[0], RegexTree) and len(self._elements[0].references):
             for string in self._elements[0].next() if index == 0 else self._elements[0].current:
-                result.append((string, {}))
+                result.append((string, {self._elements[0]: string}))
         else:
             for string in self._elements[0].next() if index == 0 else self._elements[0].current:
-                result.append((string, {self._elements[0]: string}))
+                result.append((string, {}))
 
         done = self._elements[0].done
 
         for i, element in enumerate(self._elements[1:], start=1):
             temp = []
-            if isinstance(element, CharClass):
-                for sfx in element.next() if i == index else element.current:
-                    for pfx in result:
-                        temp.append((pfx[0] + sfx, pfx[1]))
-            elif isinstance(element, RegexTree):
-                for sfx in element.next() if i == index else element.current:
-                    for pfx in result:
-                        temp.append((pfx[0] + sfx, {**pfx[1], element: sfx}))
-            else:
+            if isinstance(element, BackReference):
                 if i == index:
                     element.next()
                 for pfx in result:
@@ -146,6 +138,14 @@ def next(self) -> set[str]:
                     for sfx in element.current[reference]:
                         temp.append(
                             (pfx[0] + sfx, pfx[1]))
+            elif isinstance(element, RegexTree) and len(element.references):
+                for sfx in element.next() if i == index else element.current:
+                    for pfx in result:
+                        temp.append((pfx[0] + sfx, {**pfx[1], element: sfx}))
+            else:
+                for sfx in element.next() if i == index else element.current:
+                    for pfx in result:
+                        temp.append((pfx[0] + sfx, pfx[1]))
             result = temp
             done = done and element.done
 
@@ -204,33 +204,34 @@ def _first(self) -> set[str]:
 
         result: list[tuple[str, dict[RegexTree, str]]] = []
 
-        if isinstance(self._elements[0], CharClass):
+        if isinstance(self._elements[0], RegexTree) and len(self._elements[0].references):
             for char in self._elements[0].current:
-                result.append((char, {}))
+                result.append((char, {self._elements[0]: char}))
         else:
             for char in self._elements[0].current:
-                result.append((char, {self._elements[0]: char}))
+                result.append((char, {}))
 
         done = self._elements[0].done
 
         for element in self._elements[1:]:
             temp: list[tuple[str, dict[RegexTree, str]]] = []
             done = done and element.done
-            if isinstance(element, CharClass):
-                for pfx in result:
-                    for sfx in element.current:
-                        temp.append((pfx[0] + sfx, pfx[1]))
-            elif isinstance(element, RegexTree):
-                for pfx in result:
-                    for sfx in element.current:
-                        temp.append((pfx[0] + sfx, {**pfx[1], element: sfx}))
-            else:
+            if isinstance(element, BackReference):
                 for pfx in result:
                     reference = pfx[1][element.reference]
                     assert reference is not None
                     for sfx in element.current[reference]:
                         temp.append(
                             (pfx[0] + sfx, pfx[1]))
+            elif isinstance(element, RegexTree) and len(element.references):
+                for pfx in result:
+                    for sfx in element.current:
+                        temp.append((pfx[0] + sfx, {**pfx[1], element: sfx}))
+            else:
+                for pfx in result:
+                    for sfx in element.current:
+                        temp.append((pfx[0] + sfx, pfx[1]))
+
             result = temp
 
         self.done = done