Method get_keep_elements

miniwob++/html_tools/html_parser.py:283–349

(self, tree: html.HtmlElement, keep: list[str], max_depth: int, max_children: int, 
                            max_sibling: int, dfs_count: int=1, keep_parent: bool=False)

    # From mind2web, https://github.com/OSU-NLP-Group/Mind2Web/blob/main/src/data_utils/dom_utils.py
    def get_keep_elements(self, tree: html.HtmlElement, keep: list[str], max_depth: int, max_children: int, 
                            max_sibling: int, dfs_count: int=1, keep_parent: bool=False) -> list[str]:
        def get_anscendants(node: html.HtmlElement, max_depth: int, current_depth: int=0) -> list[str]:
            if current_depth > max_depth:
                return []

            anscendants = []
            parent = node.getparent()
            if parent is not None:
                anscendants.append(parent)
                anscendants.extend(get_anscendants(parent, max_depth, current_depth + 1))

            return anscendants
        
        def get_descendants(node: html.HtmlElement, max_depth: int, current_depth: int=0) -> list[str]:
            if current_depth > max_depth:
                return []

            descendants = []
            for child in node:
                descendants.append(child)
                descendants.extend(get_descendants(child, max_depth, current_depth + 1))

            return descendants
        
        to_keep = set(copy.deepcopy(keep))
        nodes_to_keep = set()
        
        for _ in range(max(1, dfs_count)):
            for bid in to_keep:
                candidate_node = self.get_node_by_bid(tree, bid)
                if candidate_node is None:
                    continue
                
                nodes_to_keep.add(candidate_node.attrib[self.id_attr])
                # get all ancestors or with max depth
                nodes_to_keep.update([x.attrib.get(self.id_attr, '') for x in get_anscendants(candidate_node, max_depth)])
                
                # get descendants with max depth
                nodes_to_keep.update([x.attrib.get(self.id_attr, '') for x in get_descendants(candidate_node, max_depth)][:max_children])
                # get siblings within range
                parent = candidate_node.getparent()
                if parent is None:
                    continue
                
                siblings = [x for x in parent.getchildren() if x.tag != 'text']
                if candidate_node not in siblings:
                    continue
                
                idx_in_sibling = siblings.index(candidate_node)
                nodes_to_keep.update([x.attrib.get(self.id_attr, '') 
                                    for x in siblings[max(0, idx_in_sibling - max_sibling) : idx_in_sibling + max_sibling + 1]])
            
            max_children = int(max_children * 0.5)
            max_depth = int(max_depth * 0.5)
            max_sibling = int(max_sibling * 0.7)
            
            to_keep = copy.deepcopy(nodes_to_keep)