treesearchlib 1.1.0__py3-none-any.whl

treesearch/parsers/email_html_md.py

@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+"""邮件 HTML 正文 → Markdown 转写（ADR-0005 决策 3）。
+
+body_html 存在时优先转写为 Markdown 作为邮件文档正文（保持标题/列表/表格
+结构），搜索语料与预览同源。转换复用项目已有依赖 markitdown。
+
+图片一律剥除（留 alt 占位文字）：
+- ``cid:`` 内嵌图——本期不做 cid→附件映射重写（ADR-0005 决策 4）
+- ``data:`` base64 图——避免 MB 级 base64 混进索引文本
+- 远程图——预览时加载外联图会向发件方泄漏已读回执（tracking pixel）
+"""
+import io
+import logging
+import re
+
+logger = logging.getLogger(__name__)
+
+_CONVERTER = None  # MarkItDown 懒加载单例（初始化有开销，逐邮件新建太慢）
+
+
+def _get_converter():
+    global _CONVERTER
+    if _CONVERTER is None:
+        from markitdown import MarkItDown
+
+        _CONVERTER = MarkItDown()
+    return _CONVERTER
+
+
+# markdown 图片语法：![alt](url)
+_RE_MD_IMAGE = re.compile(r"!\[([^\]]*)\]\([^)]*\)")
+
+# 转换后残留的 HTML img 标签（markitdown 对复杂嵌套偶有泄漏）
+_RE_HTML_IMAGE = re.compile(r"<img\b[^>]*>", re.IGNORECASE)
+
+_RE_BLANK_LINES = re.compile(r"\n{3,}")
+
+
+def email_html_to_md(html: str) -> str:
+    """把邮件 body_html 转写为 Markdown 文本。
+
+    Returns:
+        转写后的 Markdown；输入为空或转换失败时返回 ""（调用方退回纯文本正文）。
+    """
+    if not html or not html.strip():
+        return ""
+    try:
+        result = _get_converter().convert_stream(
+            io.BytesIO(html.encode("utf-8", errors="replace")),
+            file_extension=".html",
+        )
+        md = result.text_content or ""
+    except Exception as e:
+        logger.debug("email html->md conversion failed: %s", e)
+        return ""
+
+    # 剥图片：md 语法留 alt 占位文字；残留 HTML img 标签整段移除
+    md = _RE_MD_IMAGE.sub(lambda m: m.group(1) or "", md)
+    md = _RE_HTML_IMAGE.sub("", md)
+    return _RE_BLANK_LINES.sub("\n\n", md).strip()

treesearch/parsers/excel_parser.py

@@ -0,0 +1,218 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing(xuming624@qq.com)
+@description: Excel/Spreadsheet parser for TreeSearch.
+
+Requires optional dependency: ``pip install openpyxl``
+Extracts sheets, headers, and row data from Excel files and builds tree structure.
+"""
+import logging
+import os
+from itertools import chain
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+# Extensions supported by openpyxl
+EXCEL_EXTENSIONS = frozenset({".xlsx", ".xlsm", ".xltx", ".xltm"})
+
+
+def _non_empty_count(row) -> int:
+    """Count non-empty cells in a row tuple (None / blank 字符串视为空)."""
+    if not row:
+        return 0
+    return sum(1 for c in row if c is not None and str(c).strip())
+
+
+def _extract_excel_data(
+    excel_path: str,
+    *,
+    max_rows_per_sheet: int = 10000,
+    max_consecutive_empty_rows: int = 100,
+) -> list[dict]:
+    """Extract sheet data from an Excel file.
+
+    Returns a flat node list with:
+    - Level 1: Sheet name
+    - Level 2: Header row (columns)
+    - Level 3: Data rows
+
+    Args:
+        excel_path: Path to the Excel file.
+        max_rows_per_sheet: Maximum number of rows to process per sheet.
+        max_consecutive_empty_rows: Stop parsing a sheet after this many
+            consecutive empty rows.
+    """
+    try:
+        from openpyxl import load_workbook
+    except ImportError:
+        raise ImportError(
+            "Excel support requires 'openpyxl'. Install with: pip install openpyxl"
+        )
+
+    wb = load_workbook(excel_path, read_only=True, data_only=True)
+    nodes = []
+    row_counter = 1
+
+    for sheet_name in wb.sheetnames:
+        ws = wb[sheet_name]
+        rows_iter = ws.iter_rows(values_only=True)
+
+        # Read header row first
+        try:
+            header_row = next(rows_iter)
+        except StopIteration:
+            nodes.append({
+                "title": sheet_name,
+                "level": 1,
+                "text": "(empty sheet)",
+                "line_num": row_counter,
+                "line_start": row_counter,
+                "line_end": row_counter,
+            })
+            row_counter += 1
+            continue
+
+        # Title-row detection: many sheets start with a merged title cell
+        # (e.g., "2025年通讯录" spanning all columns) above the real headers.
+        # If row 1 has only 1 non-empty cell while row 2 has ≥2, treat row 2
+        # as the real header. Strict threshold avoids false positives on
+        # regular tables where row 1 happens to have a blank trailing cell.
+        try:
+            second_row = next(rows_iter)
+        except StopIteration:
+            second_row = None
+        first_data_row = second_row
+        if (second_row is not None
+                and _non_empty_count(header_row) <= 1
+                and _non_empty_count(second_row) >= 2):
+            # row 1 was title; row 2 becomes header; first data row unknown yet
+            header_row = second_row
+            first_data_row = None
+
+        headers = [str(cell) if cell is not None else "" for cell in header_row]
+        header_text = f"Columns: {', '.join(h for h in headers if h)}"
+
+        # Sheet node (level 1)
+        sheet_start = row_counter
+        sheet_text_parts = [header_text]
+
+        # Collect data rows with early termination.
+        # If we peeked row 2 for title detection but kept row 1 as header,
+        # row 2 is the first data row and must be prepended back to the stream.
+        data_rows_iter = (
+            chain([first_data_row], rows_iter)
+            if first_data_row is not None
+            else rows_iter
+        )
+        data_rows_text = []
+        consecutive_empty = 0
+        rows_processed = 0
+        truncated = False
+
+        for row in data_rows_iter:
+            rows_processed += 1
+            if rows_processed > max_rows_per_sheet:
+                truncated = True
+                break
+
+            cells = [str(cell) if cell is not None else "" for cell in row]
+            is_empty = not any(c.strip() for c in cells)
+
+            if is_empty:
+                consecutive_empty += 1
+                if consecutive_empty >= max_consecutive_empty_rows:
+                    truncated = True
+                    break
+                continue
+
+            # Reset counter on non-empty row
+            consecutive_empty = 0
+
+            row_text = "; ".join(
+                f"{h}: {v}" for h, v in zip(headers, cells) if v.strip()
+            )
+            if row_text:
+                data_rows_text.append(row_text)
+
+        # Combine sheet content
+        if data_rows_text:
+            # Show up to 200 rows in text to avoid excessive node size
+            displayed_rows = data_rows_text[:200]
+            sheet_text_parts.extend(displayed_rows)
+            if len(data_rows_text) > 200:
+                sheet_text_parts.append(f"... ({len(data_rows_text) - 200} more rows)")
+
+        if truncated:
+            sheet_text_parts.append(
+                f"(parsing stopped: reached limit of {max_rows_per_sheet} rows "
+                f"or {max_consecutive_empty_rows} consecutive empty rows)"
+            )
+
+        sheet_text = "\n".join(sheet_text_parts)
+        row_count = len(data_rows_text)
+
+        nodes.append({
+            "title": f"{sheet_name} ({row_count} rows)",
+            "level": 1,
+            "text": sheet_text,
+            "line_num": sheet_start,
+            "line_start": sheet_start,
+            "line_end": sheet_start + row_count,
+        })
+        row_counter += row_count + 1
+
+    wb.close()
+    return nodes
+
+
+async def excel_to_tree(
+    excel_path: str,
+    *,
+    model: Optional[str] = None,
+    if_add_node_summary: bool = True,
+    summary_chars_threshold: int = 600,
+    if_add_doc_description: bool = False,
+    if_add_node_text: bool = False,
+    if_add_node_id: bool = True,
+    **kwargs,
+) -> dict:
+    """Build a tree index from an Excel file.
+
+    Each worksheet becomes a level-1 node. Headers and row data
+    are extracted as structured text content.
+
+    Returns:
+        {'doc_name': str, 'structure': list, 'source_path': str}
+    """
+    from ..config import get_config
+
+    doc_name = os.path.splitext(os.path.basename(excel_path))[0]
+    logger.debug("Parsing Excel: %s", excel_path)
+
+    cfg = get_config()
+    nodes = _extract_excel_data(
+        excel_path,
+        max_rows_per_sheet=cfg.xlsx_max_rows_per_sheet,
+        max_consecutive_empty_rows=cfg.xlsx_max_consecutive_empty_rows,
+    )
+
+    if not nodes:
+        # Empty workbook, create a single root node
+        nodes = [{"title": doc_name, "level": 1, "text": "(empty workbook)",
+                  "line_num": 1, "line_start": 1, "line_end": 1}]
+
+    from ..indexer import _build_tree, _finalize_tree
+
+    tree = _build_tree(nodes)
+
+    return _finalize_tree(
+        tree, doc_name,
+        source_path=os.path.abspath(excel_path),
+        source_type="excel",
+        if_add_node_id=if_add_node_id,
+        if_add_node_summary=if_add_node_summary,
+        summary_chars_threshold=summary_chars_threshold,
+        if_add_node_text=if_add_node_text,
+        if_add_doc_description=if_add_doc_description,
+    )

treesearch/parsers/html_parser.py

@@ -0,0 +1,172 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing(xuming624@qq.com)
+@description: HTML parser for TreeSearch.
+
+Requires optional dependency: ``pip install beautifulsoup4``
+Extracts heading structure and text from HTML using BeautifulSoup.
+"""
+import logging
+import os
+import re
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+def _extract_html_structure(html_content: str) -> tuple[list[dict], str]:
+    """Extract headings and body text from HTML content.
+
+    Returns:
+        (headings, plain_text) where headings is a list of
+        {'title': str, 'line_num': int, 'level': int}
+    """
+    try:
+        from bs4 import BeautifulSoup
+    except ImportError:
+        raise ImportError(
+            "HTML parser requires 'beautifulsoup4'. Install with: pip install beautifulsoup4"
+        )
+
+    soup = BeautifulSoup(html_content, "html.parser")
+
+    # Remove script and style elements
+    for tag in soup(["script", "style"]):
+        tag.decompose()
+
+    # Extract plain text
+    plain_text = soup.get_text(separator="\n")
+    # Collapse blank lines
+    plain_text = re.sub(r"\n{3,}", "\n\n", plain_text).strip()
+    lines = plain_text.split("\n")
+
+    # Extract headings (h1-h6)
+    headings = []
+    heading_tags = soup.find_all(re.compile(r"^h[1-6]$"))
+    for tag in heading_tags:
+        level = int(tag.name[1])
+        title = tag.get_text(strip=True)
+        if not title:
+            continue
+        # Find approximate line number by searching in plain text
+        line_num = 1
+        title_lower = title.lower().strip()
+        for i, line in enumerate(lines):
+            if title_lower in line.lower().strip():
+                line_num = i + 1
+                break
+        headings.append({
+            "title": title,
+            "line_num": line_num,
+            "level": level,
+        })
+
+    return headings, plain_text
+
+
+async def html_to_tree(
+    html_path: str,
+    *,
+    model: Optional[str] = None,
+    if_add_node_summary: bool = True,
+    summary_chars_threshold: int = 600,
+    if_add_doc_description: bool = False,
+    if_add_node_text: bool = False,
+    if_add_node_id: bool = True,
+    **kwargs,
+) -> dict:
+    """Build a tree index from an HTML file.
+
+    Uses BeautifulSoup to extract headings (h1-h6) and text content.
+    Falls back to ``text_to_tree`` if no headings are found.
+
+    Returns:
+        {'doc_name': str, 'structure': list, 'source_path': str}
+    """
+    doc_name = os.path.splitext(os.path.basename(html_path))[0]
+    logger.debug("Parsing HTML: %s", html_path)
+
+    with open(html_path, "r", encoding="utf-8", errors="replace") as f:
+        html_content = f.read()
+
+    return await html_content_to_tree(
+        html_content,
+        doc_name=doc_name,
+        source_path=os.path.abspath(html_path),
+        model=model,
+        if_add_node_summary=if_add_node_summary,
+        summary_chars_threshold=summary_chars_threshold,
+        if_add_doc_description=if_add_doc_description,
+        if_add_node_text=if_add_node_text,
+        if_add_node_id=if_add_node_id,
+        **kwargs,
+    )
+
+
+async def html_content_to_tree(
+    html_content: str,
+    *,
+    doc_name: str,
+    source_path: str,
+    model: Optional[str] = None,
+    if_add_node_summary: bool = True,
+    summary_chars_threshold: int = 600,
+    if_add_doc_description: bool = False,
+    if_add_node_text: bool = False,
+    if_add_node_id: bool = True,
+    source_type: str = "html",
+    **kwargs,
+) -> dict:
+    """Build a tree index from HTML content in memory.
+
+    Shared core of ``html_to_tree``; also used by the MHTML parser
+    (which unpacks the HTML part from the MIME archive first).
+    """
+    headings, plain_text = _extract_html_structure(html_content)
+
+    if not headings:
+        from ..indexer import text_to_tree
+        result = await text_to_tree(
+            text_content=plain_text,
+            model=model,
+            if_add_node_summary=if_add_node_summary,
+            summary_chars_threshold=summary_chars_threshold,
+            if_add_doc_description=if_add_doc_description,
+            if_add_node_text=if_add_node_text,
+            if_add_node_id=if_add_node_id,
+            **kwargs,
+        )
+        result["doc_name"] = doc_name
+        result["source_path"] = source_path
+        return result
+
+    # Build nodes from headings
+    from ..indexer import _build_tree, _finalize_tree
+
+    lines = plain_text.split("\n")
+    nodes = []
+    for i, hd in enumerate(headings):
+        start = hd["line_num"] - 1
+        end = headings[i + 1]["line_num"] - 1 if i + 1 < len(headings) else len(lines)
+        text = "\n".join(lines[start:end]).strip()
+        nodes.append({
+            "title": hd["title"],
+            "line_num": hd["line_num"],
+            "line_start": hd["line_num"],
+            "line_end": end,
+            "level": hd["level"],
+            "text": text,
+        })
+
+    tree = _build_tree(nodes)
+
+    return _finalize_tree(
+        tree, doc_name,
+        source_path=source_path,
+        source_type=source_type,
+        if_add_node_id=if_add_node_id,
+        if_add_node_summary=if_add_node_summary,
+        summary_chars_threshold=summary_chars_threshold,
+        if_add_node_text=if_add_node_text,
+        if_add_doc_description=if_add_doc_description,
+    )