Ejemplo n.º 1
0
    def plot_tree(self,
                  filepath,
                  edge_prop_metric=None,
                  node_label_func=None,
                  node_shape_func=None):
        """
        Function that is returning the dotprogram and saving the png in filepath.

        :param filepath: Path of the file
        :param node_metric_col_print_dict: Dict of what we want to print in the node
            and in which format.
            The structure of the Dict is: {'metric_name': {'type': type, 'digits': digits}}
            type must be in ['float', 'percent', 'int']
        :param edge_prop_metric: Metric to compute the proportion on each edge
        :param node_label_func: Function to be apply to each node to get the label
        :param node_shape_func: Function to be apply to each node to get the shape
        :return: dotprogram in a string format
        """

        render_tree = \
            RenderTreeGraph(node=self.tree,
                            nodeattrfunc=lambda node:
                            self._get_node_attr(node,
                                                node_label_func,
                                                node_shape_func),
                            edgeattrfunc=lambda node, child:
                            self._get_edge_attr(node, child, edge_prop_metric))

        render_tree.to_picture(filepath)
        return self._to_dot(render_tree)
def render_tree(tree, filepath):
    """Wrapper around rendering trees into `png` or `dot` formats. Tree are 
    rendered from the root node.
    
    Args:
        tree (NodeMixin): 
        filepath (str|Path): 
    """
    from anytree.dotexport import RenderTreeGraph
    tree = RenderTreeGraph(tree.root)
    base, ext = os.path.splitext(filepath)
    if ext == '.png':
        tree.to_picture(filename=filepath)
    elif ext == '.dot':
        tree.to_dotfile(filename=filepath)
    else:
        raise Exception