yaflux._viz — yaflux documentation

Skip to main content
`Ctrl`+`K`

# Source code for yaflux._viz

    import shutil

    def _check_graphviz():
        """Check if graphviz is available"""
        try:
            import graphviz  # noqa

            return True
        except ImportError:
            return False

    def visualize_dependencies(
        self,
        fontname: str = "Helvetica",
        fontsize: int = 11,
        rankdir: str = "LR",
        complete_linecolor: str = "darkgreen",
        complete_fillcolor: str = "palegreen",
        incomplete_linecolor: str = "gray70",
        incomplete_fillcolor: str = "white",
        step_linecolor: str = "navy",
        step_fillcolor: str = "lightblue",
    ):
        """Create a clear visualization of step dependencies using Graphviz.

        Parameters
        ----------
        fontname : str, optional
            The font to use for text, by default "Helvetica"
        fontsize : int, optional
            The font size to use, by default 11
        rankdir : str, optional
            The direction of the graph layout, by default "LR"
        complete_linecolor : str, optional
            The color of edges for completed steps, by default "darkgreen"
        complete_fillcolor : str, optional
            The fill color for nodes of completed steps, by default "palegreen"
        incomplete_linecolor : str, optional
            The color of edges for incomplete steps, by default "gray70"
        incomplete_fillcolor : str, optional
            The fill color for nodes of incomplete steps, by default "white"
        step_linecolor : str, optional
            The color of edges for step nodes, by default "navy"
        step_fillcolor : str, optional
            The fill color for step nodes, by default "lightblue"

        Returns
        -------
        graphviz.Digraph
            The rendered graph object

        Examples
        --------
        >>> import yaflux as yf
        >>>
        >>> class MyAnalysis(yf.Base):
        >>>     @yf.step(creates="a")
        >>>     def step_a(self):
        >>>         return 42
        >>>
        >>>     @yf.step(creates="b", requires="a")
        >>>     def step_b(self):
        >>>         return 42
        >>>
        >>>     @yf.step(creates="c", requires=["a", "b"])
        >>>     def step_c(self):
        >>>         return 42
        >>>
        >>> analysis = MyAnalysis()
        >>>
        >>> # Visualize the dependencies
        >>> analysis.visualize_dependencies()
        >>>
        >>> # Save the visualization to a file
        >>> dot = analysis.visualize_dependencies()
        >>> dot.render('dependencies.pdf')
        """
        if not _check_graphviz():
            raise ImportError(
                "Graphviz is required for this method.\n"
                "Install with `pip install yaflux[viz]`"
            )

        # Checks if `dot` in the environment
        if not shutil.which("dot"):
            raise FileNotFoundError("Graphviz executables not found in PATH")

        from graphviz import Digraph

        dot = Digraph(comment="Analysis Dependencies")
        dot.attr(rankdir=rankdir)

        # Set some global attributes for nicer appearance
        dot.attr("node", fontname=fontname)
        dot.attr("edge", fontname=fontname)
        dot.attr("graph", fontsize=str(fontsize))

        # Track all nodes to avoid duplicates
        result_nodes = set()

        # Add all nodes and edges
        for step_name in self.available_steps:
            method = getattr(self.__class__, step_name)
            is_complete = step_name in self.completed_steps

            # Add step node
            dot.node(
                f"step_{step_name}",
                step_name,
                shape="box",
                style="filled",
                fillcolor=step_fillcolor,
                color=step_linecolor if is_complete else incomplete_linecolor,
            )

            # Add nodes for results this step creates
            for result in method.creates:
                if result not in result_nodes:
                    is_result_complete = hasattr(self.results, result)
                    color = (
                        complete_linecolor if is_result_complete else incomplete_linecolor
                    )
                    fillcolor = (
                        complete_fillcolor if is_result_complete else incomplete_fillcolor
                    )
                    dot.node(
                        result,
                        result,
                        style="filled,rounded",
                        shape="box",
                        fillcolor=fillcolor,
                        color=color,
                    )
                    result_nodes.add(result)

            # Add edges from requirements to step and step to creates
            for req in method.requires:
                if req not in result_nodes:
                    dot.node(
                        req,
                        req,
                        style="filled,rounded",
                        shape="box",
                        fillcolor=incomplete_fillcolor,
                        color=incomplete_linecolor,
                    )
                    result_nodes.add(req)

                # Edge from requirement to step
                dot.edge(
                    req,
                    f"step_{step_name}",
                    "",
                    color=complete_linecolor if is_complete else incomplete_linecolor,
                )

            # Edges from step to its outputs
            for create in method.creates:
                dot.edge(
                    f"step_{step_name}",
                    create,
                    "",
                    color=complete_linecolor if is_complete else incomplete_linecolor,
                )

        return dot