webleaf.Web

  1from .Leaf import Leaf
  2from .model.WebGraphAutoEncoder import WebGraphAutoEncoder
  3from lxml import etree
  4from lxml.cssselect import CSSSelector
  5
  6# Global variable to hold the WebGraphAutoEncoder instance
  7web_graph_auto_encoder = None
  8
  9
 10class Web:
 11    """
 12    The Web class provides an interface for parsing and interacting with HTML content.
 13    It uses the WebGraphAutoEncoder to convert the HTML structure into graph-based embeddings,
 14    allowing for extraction and comparison of elements.
 15
 16    Attributes:
 17    -----------
 18    html : str
 19        The raw HTML content passed during initialization.
 20    tree : lxml.etree.ElementTree
 21        Parsed HTML tree.
 22    features : list
 23        The encoded feature vectors for each HTML element.
 24    paths : list
 25        The XPath for each HTML element in the document.
 26    leaves : list of Leaf
 27        Leaf objects representing the elements of the HTML tree, based on their embeddings.
 28    path_leaves : dict
 29        A dictionary mapping XPaths to their corresponding Leaf objects.
 30    """
 31    def __init__(self, html: str):
 32        """
 33        Initialize the Web object by parsing the provided HTML and encoding it into embeddings.
 34
 35        Parameters:
 36        -----------
 37        html : str
 38            The HTML content to be parsed and encoded.
 39
 40        Raises:
 41        -------
 42        AssertionError if the HTML content is invalid.
 43        """
 44        global web_graph_auto_encoder
 45        self.html = html
 46        if not web_graph_auto_encoder:
 47            web_graph_auto_encoder = WebGraphAutoEncoder()
 48        self.tree = etree.ElementTree(etree.HTML(html))
 49        self.features, self.paths = web_graph_auto_encoder.extract(self.tree)
 50        self.leaves = [Leaf(feat) for feat in self.features]
 51        self.path_leaves = {path: leaf for path, leaf in zip(self.paths, self.leaves)}
 52
 53    def leaf(self, xpath: str = "", css_select: str = "") -> Leaf:
 54        """
 55        Extracts a specific element from the HTML tree based on an XPath or CSS selector
 56        and returns it as a Leaf object.
 57
 58        Parameters:
 59        -----------
 60        xpath : str, optional
 61            The XPath query to locate the element. (default is "")
 62        css_select : str, optional
 63            A CSS selector to locate the element. If provided, it is converted to XPath. (default is "")
 64
 65        Returns:
 66        --------
 67        Leaf
 68            The corresponding Leaf object for the element.
 69
 70        Raises:
 71        -------
 72        AssertionError if neither an XPath nor a CSS selector is provided, or if the element is not found.
 73        """
 74        if css_select:
 75            xpath = CSSSelector(css_select).path
 76        assert xpath, "When creating a WebLeaf please provide either a xpath or css selector."
 77        elements = self.tree.xpath(xpath)
 78        assert len(elements), f"Could not find elements at xpath [{xpath}] in html."
 79        path = self.tree.getpath(elements[0])
 80        assert path in self.path_leaves, f"The element at [{path}] was not processed by webleaf."
 81        return self.path_leaves[path]
 82
 83    def find(self, leaf: Leaf):
 84        """
 85         Finds the closest matching element in the HTML tree for the given Leaf object.
 86
 87         Parameters:
 88         -----------
 89         leaf : Leaf
 90             The Leaf object to be matched against the elements in the HTML tree.
 91
 92         Returns:
 93         --------
 94         str
 95             The XPath of the closest matching element.
 96         """
 97        return self.find_n(leaf, 1)[0]
 98
 99    def find_n(self, leaf: Leaf, n):
100        """
101        Finds the top N most similar elements to the given Leaf object,
102        based on their embeddings' distance.
103
104        Parameters:
105        -----------
106        leaf : Leaf
107            The Leaf object to compare against other elements in the tree.
108        n : int
109            The number of top similar elements to return.
110
111        Returns:
112        --------
113        list of str
114            A list of XPaths corresponding to the top N most similar elements.
115        """
116        path_dists = [(path, leaf.mdist(other)) for other, path in zip(self.leaves, self.paths)]
117        sorted_path_dists = sorted(path_dists, key=lambda path_dist: path_dist[1])
118        sorted_paths = [path for path, dist in sorted_path_dists]
119        return sorted_paths[:n]