crawler-intra-mart/extractor/html_extractor.py

from bs4 import BeautifulSoup

from utils.logger import logger


def _ensure_soup(html_or_soup):
    if hasattr(html_or_soup, "select_one") and hasattr(
        html_or_soup, "find_all"
    ):
        return html_or_soup
    return BeautifulSoup(html_or_soup, "lxml")


class HTMLExtractor:
    """
    Extract main documentation content
    from rendered HTML.
    """

    def __init__(self, selectors):
        self.selectors = selectors

    def extract_main_content(
        self,
        html
    ):
        """
        Extract the most relevant content block.
        """
        result = self.extract_with_status(html)
        return result["html"]

    def extract_with_status(self, html):
        """
        Like extract_main_content but reports whether a
        configured selector matched. Accepts either an HTML
        string or a BeautifulSoup/Tag so process_page can
        avoid re-parsing.

        The returned dict now also includes ``node`` — the
        matched BS4 element — so downstream stages (asset
        downloader, asset link rewriter) can operate directly
        on the parsed tree instead of re-parsing ``html``.
        """

        soup = _ensure_soup(html)

        self.remove_noise(soup)

        for selector in self.selectors:
            try:
                node = soup.select_one(selector)

                if node:
                    text_length = len(
                        node.get_text(strip=True)
                    )

                    if text_length < 100:
                        continue

                    logger.info(
                        f"Main content found "
                        f"with selector: {selector}"
                    )

                    return {
                        "node": node,
                        "html": str(node),
                        "matched": True,
                        "selector": selector,
                        "text_length": text_length,
                    }

            except Exception as e:
                logger.warning(
                    f"Selector failed {selector}: {e}"
                )

        logger.warning(
            "No selector matched. Using fallback extraction."
        )

        fallback = self.fallback_extract(soup)

        return {
            "node": fallback,
            "html": str(fallback),
            "matched": False,
            "selector": None,
            "text_length": len(fallback.get_text(strip=True))
            if hasattr(fallback, "get_text")
            else 0,
        }

    def remove_noise(
        self,
        soup
    ):
        """
        Remove obvious noise
        before extraction.
        """

        noise_tags = [
            "script",
            "style",
            "noscript",
            "iframe",
            "svg"
        ]

        for tag_name in noise_tags:

            for tag in soup.find_all(tag_name):

                tag.decompose()

    def fallback_extract(
        self,
        soup
    ):
        """
        Heuristic fallback extraction.
        """

        candidates = []

        possible_selectors = [
            "main",
            "article",
            ".content",
            ".document",
            ".markdown-body",
            ".theme-doc-markdown",
            ".rst-content",
            ".page-content"
        ]

        for selector in possible_selectors:

            for node in soup.select(selector):

                score = self.calculate_score(node)

                candidates.append(
                    (score, node)
                )

        if candidates:

            candidates.sort(
                key=lambda x: x[0],
                reverse=True
            )

            best = candidates[0][1]

            logger.info(
                "Fallback extraction selected "
                "best candidate."
            )

            return best

        # Final fallback = body
        body = soup.body

        if body:
            return body

        return soup

    def calculate_score(
        self,
        node
    ):
        """
        Score content relevance.
        """

        text = node.get_text(" ", strip=True)

        text_length = len(text)

        paragraph_count = len(
            node.find_all("p")
        )

        code_blocks = len(
            node.find_all(["pre", "code"])
        )

        headings = len(
            node.find_all(
                ["h1", "h2", "h3"]
            )
        )

        score = (
            text_length
            + paragraph_count * 50
            + code_blocks * 100
            + headings * 30
        )

        return score