visualize_rules_graph.py

import argparse
import os
import re
import toml
import graphviz
from dataclasses import dataclass

# REQUIRED: pip install toml
# REQUIRED: pip install graphviz


@dataclass(frozen=True, eq=True)
class Edge:
    to: str
    scope: str


def collect_rules_groups_edges():
    for config_path in args.configurations_path:
        config_path = os.path.abspath(config_path) + os.path.sep
        rules_file = config_path + 'rules.toml'

        # rules.toml should exist, deliberately fail otherwise
        rules_toml_dict = toml.load(rules_file)
        collect_rules_and_groups(rules_toml_dict)

        # edges are optional for input rules (i.e., non built-in)
        edges_file = config_path + 'edges.toml'
        if os.path.exists(edges_file):
            edges_toml_dict = toml.load(edges_file)
            collect_edges(edges_toml_dict)


def collect_rules_and_groups(rules_toml_dict):
    """
    Collects rules and groups to further build graph nodes.

    Rules belonging to a group different than `Cleanup Rule` are displayed below the group name.

    Cleanup Rule's are added to a set as they are displayed differently.
        They get an additional label `(Cleanup_Rule)`.

    Nodes without a group are added to a different set.
    """
    for rule_toml in rules_toml_dict['rules']:
        rule_name: str = sanitize_name(rule_toml['name'])
        if 'query' not in rule_toml:
            dummy_nodes.add(rule_name)

        if 'groups' in rule_toml:
            rule_groups = rule_toml['groups']
            collect_node_for_rule_with_group(rule_name, rule_groups)
        else:
            nodes_without_groups.add(rule_name)


def collect_node_for_rule_with_group(rule_name: str, rule_groups: 'list[str]'):
    for group_name in rule_groups:
        group_name = sanitize_name(group_name)
        if group_name != 'Cleanup_Rule':
            if group_name in rules_by_group_dict:
                rules_by_group_dict[group_name].append(rule_name)
            else:
                rules_by_group_dict[group_name] = [rule_name]
        else:  # we don't want to group `Cleanup Rule`s under the same graphviz shape
            cleanup_rules.add(rule_name)


def sanitize_name(s: str) -> str:
    """Graphviz does not like names with spaces. Converts spaces to '_'."""
    s = re.sub(r"\s+", '_', s)
    return s


def collect_edges(edges_toml_dict):
    """
    Groups outgoing edges by rule/group.

    All names are sanitized to replace empty spaces by `_`.
    """
    for edge_toml in edges_toml_dict['edges']:
        from_node = sanitize_name(edge_toml['from'])
        to_nodes: 'list[str]' = edge_toml['to']
        scope = sanitize_name(edge_toml['scope'])
        for to_node in to_nodes:
            to_node = sanitize_name(to_node)
            edge = Edge(to=to_node, scope=scope)
            if from_node in outgoing_edges_by_node:
                outgoing_edges_by_node[from_node].append(edge)
            else:
                outgoing_edges_by_node[from_node] = [edge]


def initialize_graph() -> graphviz.Digraph:
    graph_attr = {
        'label': str(args.title),
        'labelloc': 't',
        'fontsize': '30'
    }
    graph = graphviz.Digraph(filename=output_file_path, graph_attr=graph_attr)
    graph.format = 'svg'
    return graph


def generate_graph_nodes():
    generate_nodes_with_groups_and_outgoing_edges()
    # rules *without* outgoing edges and no groups have not been added yet
    # this is because we focus on outgoing edges when traversing `edges.toml`
    generate_nodes_without_groups_and_no_outgoing_edges()


def generate_nodes_with_groups_and_outgoing_edges():
    for rule_name in all_rule_names_with_groups_and_outgoing_edges():
        if rule_name in rules_by_group_dict:
            # several (n >= 1) rules under the same graphviz `record` shape.
            generate_node_for_group(rule_name)
        else:
            # not a graphviz `record` node. single rule in the shape.
            generate_node_for_rule_not_under_a_group(rule_name)

        added_nodes.add(rule_name)


def all_rule_names_with_groups_and_outgoing_edges() -> 'list[str]':
    # set difference
    # map.keys() is a set view, it doesn't have set's methods but supports operators
    node_names_with_only_outgoing_edges = rules_by_group_dict.keys() - \
        outgoing_edges_by_node.keys()

    rule_names_with_groups = list(outgoing_edges_by_node.keys())
    rule_names_with_groups.extend(node_names_with_only_outgoing_edges)
    rule_names_with_groups.extend(cleanup_rules)
    return rule_names_with_groups


def generate_node_for_group(rule_name: str):
    rule_names_in_group = rules_by_group_dict[rule_name]
    for group_rule_name in rule_names_in_group:
        added_nodes.add(group_rule_name)

    rule_names_label: 'list[str]' = [
        append_cleanup_rule_if_needed(rule_name) for rule_name in rule_names_in_group
    ]

    #############################
    # boolean_expression_simplify
    #
    # simplify_not_false
    # simplify_not_true
    # ...
    #############################
    node_label = rule_name + '\\n\\n' + '\\n'.join(rule_names_label)

    graph.node(rule_name, node_label, shape='record')


def append_cleanup_rule_if_needed(rule_name: str) -> str:
    """
    Should be called for rules under a group.

    If a rule is a cleanup rule, we append the label to the node's name *on the same line*.
    Cleanup Rules are treated differently because:
      1) Currently, there are no edges to `Cleanup Rule` (as opposed to other groups)
      2) nodes may have another group with incoming/outgoing edges.

    We want to display rules under a group which indicates flow (i.e., has edges).
    At the same time, we still want to indicate in the graph that a rule is a cleanup rule.
    """
    if rule_name in cleanup_rules:
        return f'{rule_name} (Cleanup Rule)'
    else:
        return rule_name


def generate_node_for_rule_not_under_a_group(rule_name: str):
    """The rule will be a standalone node; we can add (Cleanup Rule) *on a new line* if needed."""
    if rule_name not in added_nodes:
        if rule_name in dummy_nodes:
            graph.node(rule_name, shape='doubleoctagon')
        elif rule_name in cleanup_rules:
            node_label = f'{rule_name}\\n(Cleanup Rule)'
            graph.node(rule_name, node_label)
        else:
            graph.node(rule_name)


def generate_nodes_without_groups_and_no_outgoing_edges():
    # nodes that don't have any group, not even 'Cleanup_Rule'
    # and also no outgoing edge
    # i.e., leaf nodes with no groups -> standalone shape
    for node in nodes_without_groups:
        # avoid adding a node already added through an edge
        if node not in added_nodes:
            added_nodes.add(node)
            graph.node(node)


def generate_graph_edges():
    for node, edges in outgoing_edges_by_node.items():
        for edge in edges:
            graph.edge(node, edge.to, edge.scope)


###########
# Arguments
example_text = '''For a complete output, the script needs the directory with the built-in rules for a given language. Example:

        python visualize_rules_graph.py ./ff_cleanup.dot src/cleanup_rules/java demo/feature_flag_cleanup/java/configurations --title "Feature Flag Cleanup"

    Experimental:
        The generated graph may end very wide if you have several rules with no outgoing edges.
        You can experiment passing the `--unflatten` option and changing the values of `--stagger` (https://graphviz.readthedocs.io/en/stable/manual.html#unflatten).
        Another option is to manually edit the generated .dot file to include invisible edges (https://stackoverflow.com/a/11136488/1008952).
    '''

description_text = '''Script to output a .dot graph and svg image for visualizing how rules/groups are connected.
Please install the `toml` PyPi package: `python install toml`.
Please install the `graphviz` PyPi package: `python install graphviz`.

The script assumes that rules will have at most [0,2] groups. If a rule has two groups, one will be `Cleanup Rule`. The visualization will likely break if any rule has > 2 groups.'''

arg_parser = argparse.ArgumentParser(
    description=description_text, epilog=example_text, formatter_class=argparse.RawDescriptionHelpFormatter)
arg_parser.add_argument('output_file_path', type=str,
                        help="Path and file name/extension for the output file.")
arg_parser.add_argument('configurations_path', type=str, nargs='+',
                        help="One or more root directory that contains 'rules.toml' and 'edges.toml'")
arg_parser.add_argument('--title', nargs='?', default='',
                        help='Optional title for the graph')
arg_parser.add_argument('--unflatten', action='store_true', default=False)
arg_parser.add_argument('--stagger', nargs='?', default=2, const=2, type=int)
args = arg_parser.parse_args()

###########
# Execution
rules_by_group_dict: 'dict[str, list[str]]' = {}
outgoing_edges_by_node: 'dict[str, list[Edge]]' = {}

cleanup_rules: 'set[str]' = set()
nodes_without_groups: 'set[str]' = set()
# nodes without `query`
dummy_nodes: 'set[str]' = set()

output_file_path = os.path.abspath(args.output_file_path)

if os.path.isdir(output_file_path):
    raise ValueError(
        'output_file_path (first arg) should be a file, not a directory')

output_dir_path = os.path.dirname(output_file_path)
if not os.path.exists(output_dir_path):
    os.makedirs(output_dir_path)

collect_rules_groups_edges()

graph = initialize_graph()
added_nodes: 'set[str]' = set()
generate_graph_nodes()
generate_graph_edges()

if args.unflatten:
    graph.unflatten(stagger=args.stagger).render()
else:
    graph.render()