analyzer.py

"""
Module is used for analyzing link relationships
"""
from requests.exceptions import HTTPError

from ete3 import Tree, TreeStyle, TextFace, add_face_to_node
from .link import LinkNode
from .utils import multi_thread


def default_layout(node):
    node_style = TextFace(node.name, tight_text=True)
    add_face_to_node(node_style, node, column=0, position='branch-bottom')


default_style = TreeStyle()
default_style.show_leaf_name = False
default_style.layout_fn = default_layout


class LinkTree:
    """
    This is a class that represents a tree of links within TorBot. This can
    be used to build a tree, examine the number of nodes, check if a node
    exists within a tree, displaying the tree, and downloading the tree. It
    will be expanded in the future to meet further needs.

    Attributes:
        root (LinkNode): root node
        stop_depth (int): Depth of which to stop searching for links
    """
    def __init__(self, root_node, stop_depth=1):
        self._tree = build_tree(root_node, stop=stop_depth)

    def __len__(self):
        return len(self._tree)

    def __contains__(self, link):
        return self._tree.search_nodes(name=link)

    @property
    def children(self):
        return self._tree.get_children()

    def save(self, file_name, tree_style=default_style):
        """
        Saves LinkTree to file with given file_name
        Current file types supported are .png, .pdf, .svg

        Args:
            file_name (str): Name of file being saved to
            tree_style (TreeStyle): Styling of downloaded tree
        """
        self._tree.render(file_name, tree_style)

    def show(self, tree_style=default_style):
        """
        Displays image of LinkTree

        Args:
            tree_style (TreeStyle): Styling of downloaded tree
        """
        self._tree.show(tree_style)


def build_tree(link, stop=1, rec=0):
    """
    Builds link tree by traversing through children nodes.

    Args:
        link (LinkNode): root node of tree
        stop (int): depth of tree
        rec (int): level of recursion

    Returns:
        tree (ete3.Tree): Built tree.
    """

    tree = Tree(name=link.name)

    if rec == stop:
        return tree
    else:
        rec += 1

    for child in link.links:
        try:
            node = LinkNode(child)
        except Exception as error:
            print(f"Failed to create LinkNode for link: {child}.")
            print(f"Error: {error}")
            continue
        if node.links:
            tree.add_child(build_tree(node, stop, rec))
        else:
            tree.add_child(Tree(name=node.name))

    return tree