direct-corpus-information / dci_engine.py
Marek4321's picture
initial commit
999aff1
Raw
History Blame Contribute Delete
2.66 kB
import re
def grep_search(
lines: list[str],
keyword: str,
case_sensitive: bool = False,
context_lines: int = 3,
) -> list[dict]:
flags = 0 if case_sensitive else re.IGNORECASE
pattern = re.compile(re.escape(keyword), flags)
results = []
for idx, line in enumerate(lines):
if pattern.search(line):
ctx_start = max(0, idx - context_lines)
ctx_end = min(len(lines), idx + context_lines + 1)
context = lines[ctx_start:ctx_end]
results.append(
{
"line_number": idx + 1,
"matched_line": line,
"context": context,
"context_start": ctx_start + 1,
"keyword": keyword,
}
)
return results
def chain_search(
lines: list[str],
keywords: list[str],
context_lines: int = 3,
) -> list[dict]:
if not keywords:
return []
results = grep_search(lines, keywords[0], context_lines=context_lines)
for kw in keywords[1:]:
pattern = re.compile(re.escape(kw), re.IGNORECASE)
results = [r for r in results if pattern.search(r["matched_line"])]
# Update keyword field to reflect all keywords in the chain
for r in results:
r["keyword"] = keywords[0]
return results
def search_corpus(
corpus: dict[str, list[str]],
chain_queries: list[list[str]],
max_hits_per_query: int = 10,
) -> list[dict]:
seen: set[tuple] = set()
all_hits: list[dict] = []
def _collect(query: list[str], hits_from_search: list[dict], filename: str) -> None:
for hit in hits_from_search[:max_hits_per_query]:
key = (filename, hit["line_number"])
if key in seen:
continue
seen.add(key)
all_hits.append(
{
"file": filename,
"line_number": hit["line_number"],
"matched_line": hit["matched_line"],
"context": hit["context"],
"context_start": hit["context_start"],
"keywords": query,
}
)
for query in chain_queries:
if not query:
continue
for filename, lines in corpus.items():
hits = chain_search(lines, query)
if hits:
_collect(query, hits, filename)
else:
# Fallback: search each keyword separately when no same-line matches
for kw in query:
_collect([kw], grep_search(lines, kw), filename)
return all_hits