13379554081

Committed 17 Feb 2025 10:31PM UTC coverage: 92.989% (+0.07%) from 92.921%

Build # 13379554081

Build Type

push

github

Committed by

david-yz-liu

Commit Message

Export generate_cfg from python_ta.cfg

Run Details

1 of 1 new or added line in 1 file covered. (100.0%)

2 existing lines in 1 file now uncovered.

3263 of 3509 relevant lines covered (92.99%)

17.67 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

95.92

/python_ta/cfg/cfg_generator.py

"""
Provides a function to generate and display the control flow graph of a given module.
"""

from __future__ import annotations

import html
import importlib.util
import os.path
import sys
from typing import Any, Optional

import graphviz
from astroid import nodes
from astroid.builder import AstroidBuilder

try:
    from ..transforms.z3_visitor import Z3Visitor

    z3_dependency_available = True
except ImportError:
    Z3Visitor = Any
    z3_dependency_available = False

from .graph import CFGBlock, ControlFlowGraph
from .visitor import CFGVisitor

GRAPH_OPTIONS = {"format": "svg", "node_attr": {"shape": "box", "fontname": "Courier New"}}
SUBGRAPH_OPTIONS = {"fontname": "Courier New"}


def generate_cfg(
    mod: str = "", auto_open: bool = False, visitor_options: Optional[dict[str, Any]] = None
) -> None:
    """Generate a control flow graph for the given module.

    Supported Options:
      - "separate-condition-blocks": bool
            This option specifies whether the test condition of an if statement gets merged with any
            preceding statements or placed in a new block. By default, it will merge them.
      - "functions": list[str]
            This option specifies whether to restrict the creation of cfgs to just top-level
            function definitions or methods provided in this list. By default, it will create the
            cfg for the entire file.

    Args:
        mod (str): The path to the module. `mod` can either be the path of a file (must have `.py`
            extension) or have no argument (generates a CFG for the Python file from which this
            function is called).
        auto_open (bool): Automatically open the graph in your browser.
        visitor_options (dict): An options dict to configure how the cfgs are generated.
    """
    _generate(mod=mod, auto_open=auto_open, visitor_options=visitor_options)


def _generate(
    mod: str = "", auto_open: bool = False, visitor_options: Optional[dict[str, Any]] = None
) -> None:
    """Generate a control flow graph for the given module.

    `mod` can either be:
      - the path of a file (must have `.py` extension).
      - no argument -- generate a CFG for the Python file from which this function is called.
    """
    # Generate a control flow graph for the given file
    abs_path = _get_valid_file_path(mod)
    # Print an error message if the file is not valid and early return
    if abs_path is None:  # _get_valid_file_path returns None in case of invalid file
        return

    file_name = os.path.splitext(os.path.basename(abs_path))[0]
    module = AstroidBuilder().file_build(abs_path)

    # invoke Z3Visitor if z3 dependency is available
    if z3_dependency_available:
        z3v = Z3Visitor()
        module = z3v.visitor.visit(module)

    visitor = CFGVisitor(options=visitor_options)
    module.accept(visitor)

    _display(visitor.cfgs, file_name, auto_open=auto_open)


def _get_valid_file_path(mod: str = "") -> Optional[str]:
    """Return the valid absolute path of `mod`, a path to the target file."""
    # Allow call to check with empty args
    if mod == "":
        m = sys.modules["__main__"]
        spec = importlib.util.spec_from_file_location(m.__name__, m.__file__)
        mod = spec.origin
    # Enforce the API to only except `mod` type as str
    elif not isinstance(mod, str):
        print(
            "No CFG generated. Input to check, `{}`, has invalid type, must be a string.".format(
                mod
            )
        )
        return

    # At this point, `mod` is of type str
    if not os.path.isfile(mod):
        # `mod` is not a file so print an error message
        print("Could not find the file called, `{}`\n".format(mod))
        return

    # `mod` may be a relative path to a valid file so return its absolute path
    return os.path.abspath(mod)


def _display(
    cfgs: dict[nodes.NodeNG, ControlFlowGraph], filename: str, auto_open: bool = False
) -> None:
    graph = graphviz.Digraph(name=filename + ".gv", **GRAPH_OPTIONS)
    for node, cfg in cfgs.items():
        if isinstance(node, nodes.Module):
            subgraph_label = "__main__"
        elif isinstance(node, nodes.FunctionDef):
            scope_parent = node.scope().parent
            subgraph_label = node.name
            # Update the label to the qualified name if it is a method
            if isinstance(scope_parent, nodes.ClassDef):
                subgraph_label = scope_parent.name + "." + subgraph_label
        else:
            continue
        with graph.subgraph(name=f"cluster_{cfg.cfg_id}") as c:
            visited = set()
            _visit(cfg.start, c, visited, cfg.end)
            for block in cfg.unreachable_blocks:
                _visit(block, c, visited, cfg.end)
            c.attr(label=subgraph_label, **SUBGRAPH_OPTIONS)

    graph.render(outfile=filename + ".svg", view=auto_open)


def _visit(block: CFGBlock, graph: graphviz.Digraph, visited: set[int], end: CFGBlock) -> None:
    """
    Visit a CFGBlock and add it to the control flow graph.
    """
    node_id = f"{graph.name}_{block.id}"
    if node_id in visited:
        return

    label = ""
    fill_color = "white"

    # Identify special cases
    if len(block.statements) == 1:
        stmt = block.statements[0]
        if isinstance(stmt, nodes.Arguments):
            label = f"{stmt.as_string()}\n"
            fill_color = "palegreen"
        elif isinstance(stmt.parent, nodes.If) and stmt is stmt.parent.test:
            label = f"< if<U><B>{html.escape(stmt.as_string())}</B></U><BR/> >"
        elif isinstance(stmt.parent, nodes.While) and stmt is stmt.parent.test:
            label = f"< while<U><B>{html.escape(stmt.as_string())}</B></U><BR/> >"
        elif isinstance(stmt.parent, nodes.For) and stmt is stmt.parent.iter:
            label = f"< for {html.escape(stmt.parent.target.as_string())} in<U><B>{html.escape(stmt.as_string())}</B></U><BR/> >"
        elif isinstance(stmt.parent, nodes.For) and stmt is stmt.parent.target:
            label = f"< for<U><B>{html.escape(stmt.as_string())} </B></U> in {html.escape(stmt.parent.iter.as_string())}<BR/> >"

    if not label:  # Default
        label = "\n".join([s.as_string() for s in block.statements]) + "\n"

    # Need to escape backslashes explicitly.
    label = label.replace("\\", "\\\\")
    # \l is used for left alignment.
    label = label.replace("\n", "\\l")

    # Change the fill colour if block is the end of the cfg or unreachable
    if block == end:
        fill_color = "black"
    elif not block.reachable:
        fill_color = "grey93"

    graph.node(node_id, label=label, fillcolor=fill_color, style="filled")
    visited.add(node_id)

    for edge in block.successors:
        color = "black" if edge.is_feasible else "lightgrey"
        if edge.get_label() is not None:
            graph.edge(
                node_id, f"{graph.name}_{edge.target.id}", label=edge.get_label(), color=color
            )
        else:
            graph.edge(node_id, f"{graph.name}_{edge.target.id}", color=color)
        _visit(edge.target, graph, visited, end)

1	"""
2	Provides a function to generate and display the control flow graph of a given module.
3	"""
4
5	from __future__ import annotations	20✔
6
7	import html	20✔
8	import importlib.util	20✔
9	import os.path	20✔
10	import sys	20✔
11	from typing import Any, Optional	20✔
12
13	import graphviz	20✔
14	from astroid import nodes	20✔
15	from astroid.builder import AstroidBuilder	20✔
16
17	try:	20✔
18	from ..transforms.z3_visitor import Z3Visitor	20✔
19
20	z3_dependency_available = True	10✔
21	except ImportError:	10✔
22	Z3Visitor = Any	10✔
23	z3_dependency_available = False	10✔
24
25	from .graph import CFGBlock, ControlFlowGraph	20✔
26	from .visitor import CFGVisitor	20✔
27
28	GRAPH_OPTIONS = {"format": "svg", "node_attr": {"shape": "box", "fontname": "Courier New"}}	20✔
29	SUBGRAPH_OPTIONS = {"fontname": "Courier New"}	20✔
30
31
32	def generate_cfg(	20✔
33	mod: str = "", auto_open: bool = False, visitor_options: Optional[dict[str, Any]] = None
34	) -> None:
35	"""Generate a control flow graph for the given module.
36
37	Supported Options:
38	- "separate-condition-blocks": bool
39	This option specifies whether the test condition of an if statement gets merged with any
40	preceding statements or placed in a new block. By default, it will merge them.
41	- "functions": list[str]
42	This option specifies whether to restrict the creation of cfgs to just top-level
43	function definitions or methods provided in this list. By default, it will create the
44	cfg for the entire file.
45
46	Args:
47	mod (str): The path to the module. `mod` can either be the path of a file (must have `.py`
48	extension) or have no argument (generates a CFG for the Python file from which this
49	function is called).
50	auto_open (bool): Automatically open the graph in your browser.
51	visitor_options (dict): An options dict to configure how the cfgs are generated.
52	"""
53	_generate(mod=mod, auto_open=auto_open, visitor_options=visitor_options)	20✔
54
55
56	def _generate(	20✔
57	mod: str = "", auto_open: bool = False, visitor_options: Optional[dict[str, Any]] = None
58	) -> None:
59	"""Generate a control flow graph for the given module.
60
61	`mod` can either be:
62	- the path of a file (must have `.py` extension).
63	- no argument -- generate a CFG for the Python file from which this function is called.
64	"""
65	# Generate a control flow graph for the given file
66	abs_path = _get_valid_file_path(mod)	20✔
67	# Print an error message if the file is not valid and early return
68	if abs_path is None: # _get_valid_file_path returns None in case of invalid file	20✔
69	return	20✔
70
71	file_name = os.path.splitext(os.path.basename(abs_path))[0]	20✔
72	module = AstroidBuilder().file_build(abs_path)	20✔
73
74	# invoke Z3Visitor if z3 dependency is available
75	if z3_dependency_available:	20✔
76	z3v = Z3Visitor()	10✔
77	module = z3v.visitor.visit(module)	10✔
78
79	visitor = CFGVisitor(options=visitor_options)	20✔
80	module.accept(visitor)	20✔
81
82	_display(visitor.cfgs, file_name, auto_open=auto_open)	20✔
83
84
85	def _get_valid_file_path(mod: str = "") -> Optional[str]:	20✔
86	"""Return the valid absolute path of `mod`, a path to the target file."""
87	# Allow call to check with empty args
88	if mod == "":	20✔
89	m = sys.modules["__main__"]	×
90	spec = importlib.util.spec_from_file_location(m.__name__, m.__file__)	×
UNCOV 91	mod = spec.origin	×
92	# Enforce the API to only except `mod` type as str
93	elif not isinstance(mod, str):	20✔
94	print(	20✔
95	"No CFG generated. Input to check, `{}`, has invalid type, must be a string.".format(
96	mod
97	)
98	)
99	return	20✔
100
101	# At this point, `mod` is of type str
102	if not os.path.isfile(mod):	20✔
103	# `mod` is not a file so print an error message
104	print("Could not find the file called, `{}`\n".format(mod))	20✔
105	return	20✔
106
107	# `mod` may be a relative path to a valid file so return its absolute path
108	return os.path.abspath(mod)	20✔
109
110
111	def _display(	20✔
112	cfgs: dict[nodes.NodeNG, ControlFlowGraph], filename: str, auto_open: bool = False
113	) -> None:
114	graph = graphviz.Digraph(name=filename + ".gv", **GRAPH_OPTIONS)	20✔
115	for node, cfg in cfgs.items():	20✔
116	if isinstance(node, nodes.Module):	20✔
117	subgraph_label = "__main__"	20✔
118	elif isinstance(node, nodes.FunctionDef):	20✔
119	scope_parent = node.scope().parent	20✔
120	subgraph_label = node.name	20✔
121	# Update the label to the qualified name if it is a method
122	if isinstance(scope_parent, nodes.ClassDef):	20✔
123	subgraph_label = scope_parent.name + "." + subgraph_label	20✔
124	else:
UNCOV 125	continue	×
126	with graph.subgraph(name=f"cluster_{cfg.cfg_id}") as c:	20✔
127	visited = set()	20✔
128	_visit(cfg.start, c, visited, cfg.end)	20✔
129	for block in cfg.unreachable_blocks:	20✔
130	_visit(block, c, visited, cfg.end)	20✔
131	c.attr(label=subgraph_label, **SUBGRAPH_OPTIONS)	20✔
132
133	graph.render(outfile=filename + ".svg", view=auto_open)	20✔
134
135
136	def _visit(block: CFGBlock, graph: graphviz.Digraph, visited: set[int], end: CFGBlock) -> None:	20✔
137	"""
138	Visit a CFGBlock and add it to the control flow graph.
139	"""
140	node_id = f"{graph.name}_{block.id}"	20✔
141	if node_id in visited:	20✔
142	return	20✔
143
144	label = ""	20✔
145	fill_color = "white"	20✔
146
147	# Identify special cases
148	if len(block.statements) == 1:	20✔
149	stmt = block.statements[0]	20✔
150	if isinstance(stmt, nodes.Arguments):	20✔
151	label = f"{stmt.as_string()}\n"	20✔
152	fill_color = "palegreen"	20✔
153	elif isinstance(stmt.parent, nodes.If) and stmt is stmt.parent.test:	20✔
154	label = f"< if<U><B>{html.escape(stmt.as_string())}</B></U><BR/> >"	20✔
155	elif isinstance(stmt.parent, nodes.While) and stmt is stmt.parent.test:	20✔
156	label = f"< while<U><B>{html.escape(stmt.as_string())}</B></U><BR/> >"	20✔
157	elif isinstance(stmt.parent, nodes.For) and stmt is stmt.parent.iter:	20✔
158	label = f"< for {html.escape(stmt.parent.target.as_string())} in<U><B>{html.escape(stmt.as_string())}</B></U><BR/> >"	20✔
159	elif isinstance(stmt.parent, nodes.For) and stmt is stmt.parent.target:	20✔
160	label = f"< for<U><B>{html.escape(stmt.as_string())} </B></U> in {html.escape(stmt.parent.iter.as_string())}<BR/> >"	20✔
161
162	if not label: # Default	20✔
163	label = "\n".join([s.as_string() for s in block.statements]) + "\n"	20✔
164
165	# Need to escape backslashes explicitly.
166	label = label.replace("\\", "\\\\")	20✔
167	# \l is used for left alignment.
168	label = label.replace("\n", "\\l")	20✔
169
170	# Change the fill colour if block is the end of the cfg or unreachable
171	if block == end:	20✔
172	fill_color = "black"	20✔
173	elif not block.reachable:	20✔
174	fill_color = "grey93"	20✔
175
176	graph.node(node_id, label=label, fillcolor=fill_color, style="filled")	20✔
177	visited.add(node_id)	20✔
178
179	for edge in block.successors:	20✔
180	color = "black" if edge.is_feasible else "lightgrey"	20✔
181	if edge.get_label() is not None:	20✔
182	graph.edge(	20✔
183	node_id, f"{graph.name}_{edge.target.id}", label=edge.get_label(), color=color
184	)
185	else:
186	graph.edge(node_id, f"{graph.name}_{edge.target.id}", color=color)	20✔
187	_visit(edge.target, graph, visited, end)	20✔

pyta-uoft / pyta / 13379554081

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous