| 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"])] |
|
|
| |
| 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: |
| |
| for kw in query: |
| _collect([kw], grep_search(lines, kw), filename) |
|
|
| return all_hits |
|
|