Skip to content

Commit 5fc9fe2

Browse files
tkarabelatkarabela
committed
fix: Proper search support
1 parent d8e22aa commit 5fc9fe2

4 files changed

Lines changed: 100 additions & 57 deletions

File tree

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ Finite automaton accepting `(foo)*bar|baz`:
4343
## Features compared to standard `re` module
4444

4545
- Library
46-
- `match()`, `fullmatch()` and `search()` methods (search is currently implemented naively via match)
46+
- `match()`, `fullmatch()`, `search()` and `finditer()` methods
4747
- `Match` object containing span and matched text (but no groups)
4848
- flags `DOTALL`, `IGNORECASE` and `MULTILINE`
4949

‎src/regex_automata/regex/nfa_evaluator.py‎

Lines changed: 70 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,68 +1,94 @@
1+
from typing import Iterator
2+
13
from regex_automata.automata.nfa import NFA
24
from regex_automata.regex.flags import PatternFlag
35
from regex_automata.regex.match import Match
46

57

68
class NFAEvaluator:
9+
class Head:
10+
def __init__(self, evaluator: "NFAEvaluator", start: int) -> None:
11+
self.start = start
12+
self.evaluator = evaluator
13+
self.states: set[int] = set(self.evaluator.initial_states)
14+
self.entered_final = bool(self.states & self.evaluator.final_states)
15+
self.left_final = False
16+
17+
def step_epsilon(self, c_previous: int, c_next: int) -> None:
18+
self.states = new_states = self._step_epsilon(c_previous, c_next, self.states)
19+
new_in_final = bool(new_states & self.evaluator.final_states)
20+
self.entered_final = self.entered_final or new_in_final
21+
self.left_final = self.entered_final and not new_in_final
22+
23+
def _step_epsilon(self, c_previous: int, c_next: int, states: set[int]) -> set[int]:
24+
return self.evaluator.nfa.epsilon_closure(states, c_previous, c_next)
25+
26+
def step_read(self, c_previous: int, c_next: int) -> None:
27+
self.states = new_states = self._step_read(c_previous, c_next, self.states)
28+
new_in_final = bool(new_states & self.evaluator.final_states)
29+
self.entered_final = self.entered_final or new_in_final
30+
self.left_final = self.entered_final and not new_in_final
31+
32+
def _step_read(self, c_previous: int, c_next: int, states: set[int]) -> set[int]:
33+
assert c_next != -1
34+
new_states = set()
35+
for u in states:
36+
u_transitions = self.evaluator.nfa.transitions.get(u, {})
37+
for p, vs in u_transitions.items():
38+
if p.consume_char and p.matches(c_previous, c_next):
39+
new_states.update(vs)
40+
41+
return new_states
42+
43+
def __repr__(self) -> str:
44+
return f"<Head {self.start=} {self.states=} {self.entered_final=} {self.left_final=}>"
45+
746
def __init__(self, nfa: NFA, flags: PatternFlag = PatternFlag.NOFLAG) -> None:
847
self.nfa = nfa
9-
self.states: set[int] = self.nfa.trivial_epsilon_closure({nfa.initial_state})
10-
self.flags = flags
1148
self.initial_states = self.nfa.trivial_epsilon_closure({self.nfa.initial_state})
49+
self.heads: list["NFAEvaluator.Head"] = []
50+
self.flags = flags
1251
self.final_states = set(self.nfa.final_states)
1352

14-
def match(self, text: str, start: int = 0, end: int | None = None) -> Match | None:
53+
def finditer(self, text: str, start: int = 0, end: int | None = None, search: bool = True) -> Iterator[Match]:
1554
if self.flags & PatternFlag.IGNORECASE:
1655
text = text.lower()
1756

1857
end_ = end if end is not None else len(text)
1958

20-
entered_final = bool(self.states & self.final_states)
21-
left_final = False
59+
self.heads.append(self.Head(self, min(len(text), start)))
2260

2361
c_previous = -1
24-
for i in range(min(len(text), start), min(len(text), end_)):
62+
for char_no, i in enumerate(range(min(len(text), start), min(len(text), end_))):
63+
match_at_position = False
64+
if search and char_no > 0:
65+
self.heads.append(self.Head(self, i))
66+
2567
c_next = ord(text[i])
2668

27-
new_states = self.step_epsilon(c_previous, c_next, self.states)
28-
new_in_final = bool(new_states & self.final_states)
29-
entered_final = entered_final or new_in_final
30-
left_final = entered_final and not new_in_final
31-
if left_final:
32-
return Match.from_span_and_text(start, i, text)
33-
self.states = new_states
34-
35-
new_states = self.step_read(c_previous, c_next, self.states)
36-
new_in_final = bool(new_states & self.final_states)
37-
entered_final = entered_final or new_in_final
38-
left_final = entered_final and not new_in_final
39-
if left_final:
40-
return Match.from_span_and_text(start, i, text)
41-
self.states = new_states
69+
for head in self.heads:
70+
head.step_epsilon(c_previous, c_next)
71+
if not match_at_position and head.left_final:
72+
self.purge_heads(i-1)
73+
yield Match.from_span_and_text(head.start, i-1, text)
74+
match_at_position = True # avoid returning multiple matches
75+
76+
for head in self.heads:
77+
head.step_read(c_previous, c_next)
78+
if not match_at_position and head.left_final:
79+
self.purge_heads(i)
80+
yield Match.from_span_and_text(head.start, i, text)
81+
match_at_position = True
4282

4383
c_previous = c_next
4484

4585
c_next = -1
46-
new_states = self.step_epsilon(c_previous, c_next, self.states)
47-
new_in_final = bool(new_states & self.final_states)
48-
entered_final = entered_final or new_in_final
49-
left_final = entered_final and not new_in_final
50-
51-
if entered_final and not left_final:
52-
return Match.from_span_and_text(start, end_, text)
53-
else:
54-
return None
55-
56-
def step_epsilon(self, c_previous: int, c_next: int, states: set[int]) -> set[int]:
57-
return self.nfa.epsilon_closure(states, c_previous, c_next)
58-
59-
def step_read(self, c_previous: int, c_next: int, states: set[int]) -> set[int]:
60-
assert c_next != -1
61-
new_states = set()
62-
for u in states:
63-
u_transitions = self.nfa.transitions.get(u, {})
64-
for p, vs in u_transitions.items():
65-
if p.consume_char and p.matches(c_previous, c_next):
66-
new_states.update(vs)
67-
68-
return new_states
86+
for head in self.heads:
87+
head.step_epsilon(c_previous, c_next)
88+
89+
if head.entered_final and not head.left_final:
90+
yield Match.from_span_and_text(head.start, end_, text)
91+
return
92+
93+
def purge_heads(self, start_min: int) -> None:
94+
self.heads = [h for h in self.heads if h.start >= start_min]

‎src/regex_automata/regex/pattern.py‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
from typing import Iterator
2+
13
from .flags import PatternFlag
24
from .match import Match
35
from regex_automata.regex.nfa_evaluator import NFAEvaluator
@@ -61,13 +63,17 @@ def fullmatch(self, text: str, start: int = 0, end: int | None = None) -> Match
6163

6264
def match(self, text: str, start: int = 0, end: int | None = None) -> Match | None:
6365
evaluator = NFAEvaluator(self.nfa, self.flags)
64-
return evaluator.match(text, start, end)
66+
try:
67+
return next(evaluator.finditer(text, start, end, search=False))
68+
except StopIteration:
69+
return None
6570

6671
def search(self, text: str, start: int = 0, end: int | None = None) -> Match | None:
67-
# TODO implement this properly via automaton
68-
end_ = end if end is not None else len(text) + 1
69-
for i in range(start, end_):
70-
m = self.match(text, i, end)
71-
if m is not None:
72-
return m
73-
return None
72+
try:
73+
return next(self.finditer(text, start, end))
74+
except StopIteration:
75+
return None
76+
77+
def finditer(self, text: str, start: int = 0, end: int | None = None) -> Iterator[Match]:
78+
evaluator = NFAEvaluator(self.nfa, self.flags)
79+
yield from evaluator.finditer(text, start, end)

‎tests/test_regex.py‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import pytest
22

33
import regex_automata
4+
from regex_automata.regex.match import Match
45

56

67
@pytest.mark.parametrize("pattern,s,result",
@@ -110,6 +111,16 @@ def test_search():
110111
m = p1.search("text abc@def.com xyz@123.com", start=10)
111112
assert m is not None and m.match == "xyz@123.com"
112113

114+
115+
def test_overlapping_search():
116+
p1 = regex_automata.compile(r"aa")
117+
assert list(p1.finditer("aaaaaaa")) == [
118+
Match((0, 2), "aa"),
119+
Match((2, 4), "aa"),
120+
Match((4, 6), "aa"),
121+
]
122+
123+
113124
def test_boundary_assertion():
114125
m = regex_automata.search(r"abc$", "foo abc")
115126
assert m is not None and m.match == "abc"
@@ -120,8 +131,8 @@ def test_boundary_assertion():
120131

121132
m = regex_automata.search(r"^abc", "abc foo")
122133
assert m is not None and m.match == "abc"
123-
# m = regex_automata.search(r"^abc", "foo abc")
124-
# assert m is None
134+
m = regex_automata.search(r"^abc", "foo abc")
135+
assert m is None
125136
m = regex_automata.search(r"^abc", "foo\nabc", regex_automata.MULTILINE)
126137
assert m is not None and m.match == "abc"
127138

@@ -130,7 +141,7 @@ def test_boundary_assertion():
130141
m = regex_automata.search(r"oon\b", "moon")
131142
assert m is not None and m.span == (1, 4)
132143

133-
# m = regex_automata.search(r"\Bon", "at noon")
134-
# assert m is not None and m.span == (5, 8)
144+
m = regex_automata.search(r"\Bon", "at noon")
145+
assert m is not None and m.span == (5, 7)
135146
m = regex_automata.search(r"\Bno", "at noon")
136147
assert m is None

0 commit comments

Comments
 (0)