36646375051

Committed 29 Sep 2026 11:39PM UTC coverage: 90.867% (-0.06%) from 90.925%

Build # 36646375051

Build Type

push

github

Committed by

web-flow

Commit Message

Added type annotations and fixed mypy errors in python-ta package (#1397)

Coverage Stats

256 of 269 new or added lines in 33 files covered. (95.17%)

2 existing lines in 2 files now uncovered.

3761 of 4139 relevant lines covered (90.87%)

17.67 hits per line

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

92.16

/packages/python-ta/src/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 logging
import os.path
import sys
from typing import TYPE_CHECKING, Any, Optional

import graphviz
from astroid import nodes
from astroid.manager import AstroidManager

from .visitor import CFGVisitor

if TYPE_CHECKING:
    from python_ta_z3 import Z3Visitor

    from .graph import CFGBlock, ControlFlowGraph

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,
    z3_enabled: bool = False,
) -> 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.
        z3_enabled (bool): An option that enables z3 when True (by default False).
    """
    _generate(mod=mod, auto_open=auto_open, visitor_options=visitor_options, z3_enabled=z3_enabled)


def _generate(
    mod: str = "",
    auto_open: bool = False,
    visitor_options: Optional[dict[str, Any]] = None,
    z3_enabled: bool = False,
) -> 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 = AstroidManager().ast_from_file(abs_path)

    # invoke Z3Visitor if z3 dependency is enabled
    if z3_enabled:
        try:
            from python_ta_z3 import Z3Visitor

        except ImportError:
            logging.error("Failed to import Z3Visitor. Aborting.")
            raise
        z3v = Z3Visitor()
        module = z3v.visitor.visit(module)

    visitor = CFGVisitor(options=visitor_options, z3_enabled=z3_enabled)
    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__)
        if spec is None or spec.origin is None:
            return None
        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 None

    # 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 None

    # `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[str] = 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[str], 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 block.statements and isinstance(block.statements[0], nodes.Pattern):
        label = f"case {html.escape(block.statements[0].as_string())}"
        label += f" if {block.statements[1].as_string()}" if len(block.statements) == 2 else ""

    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 logging	20 ✔
10	import os.path	20 ✔
11	import sys	20 ✔
12	from typing import TYPE_CHECKING, Any, Optional	20 ✔
13
14	import graphviz	20 ✔
15	from astroid import nodes	20 ✔
16	from astroid.manager import AstroidManager	20 ✔
17
18	from .visitor import CFGVisitor	20 ✔
19
20	if TYPE_CHECKING:
21	from python_ta_z3 import Z3Visitor
22
23	from .graph import CFGBlock, ControlFlowGraph
24
25	GRAPH_OPTIONS = {"format": "svg", "node_attr": {"shape": "box", "fontname": "Courier New"}}	20 ✔
26	SUBGRAPH_OPTIONS = {"fontname": "Courier New"}	20 ✔
27
28
29	def generate_cfg(	20 ✔
30	mod: str = "",
31	auto_open: bool = False,
32	visitor_options: Optional[dict[str, Any]] = None,
33	z3_enabled: bool = False,
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	z3_enabled (bool): An option that enables z3 when True (by default False).
53	"""
54	_generate(mod=mod, auto_open=auto_open, visitor_options=visitor_options, z3_enabled=z3_enabled)	20 ✔
55
56
57	def _generate(	20 ✔
58	mod: str = "",
59	auto_open: bool = False,
60	visitor_options: Optional[dict[str, Any]] = None,
61	z3_enabled: bool = False,
62	) -> None:
63	"""Generate a control flow graph for the given module.
64
65	`mod` can either be:
66	- the path of a file (must have `.py` extension).
67	- no argument -- generate a CFG for the Python file from which this function is called.
68	"""
69	# Generate a control flow graph for the given file
70	abs_path = _get_valid_file_path(mod)	20 ✔
71	# Print an error message if the file is not valid and early return
72	if abs_path is None: # _get_valid_file_path returns None in case of invalid file	20 ✔
73	return	20 ✔
74
75	file_name = os.path.splitext(os.path.basename(abs_path))[0]	20 ✔
76	module = AstroidManager().ast_from_file(abs_path)	20 ✔
77
78	# invoke Z3Visitor if z3 dependency is enabled
79	if z3_enabled:	20 ✔
80	try:	10 ✔
81	from python_ta_z3 import Z3Visitor	10 ✔
82
83	except ImportError:	10 ✔
84	logging.error("Failed to import Z3Visitor. Aborting.")	10 ✔
85	raise	10 ✔
86	z3v = Z3Visitor()	10 ✔
87	module = z3v.visitor.visit(module)	10 ✔
88
89	visitor = CFGVisitor(options=visitor_options, z3_enabled=z3_enabled)	20 ✔
90	module.accept(visitor)	20 ✔
91
92	_display(visitor.cfgs, file_name, auto_open=auto_open)	20 ✔
93
94
95	def _get_valid_file_path(mod: str = "") -> Optional[str]:	20 ✔
96	"""Return the valid absolute path of `mod`, a path to the target file."""
97	# Allow call to check with empty args
98	if mod == "":	20 ✔
99	m = sys.modules["__main__"]	×
100	spec = importlib.util.spec_from_file_location(m.__name__, m.__file__)	×
NEW 101	if spec is None or spec.origin is None:	×
NEW 102	return None	×
UNCOV 103	mod = spec.origin	×
104	# Enforce the API to only except `mod` type as str
105	elif not isinstance(mod, str):	20 ✔
106	print(	20 ✔
107	"No CFG generated. Input to check, `{}`, has invalid type, must be a string.".format(
108	mod
109	)
110	)
111	return None	20 ✔
112
113	# At this point, `mod` is of type str
114	if not os.path.isfile(mod):	20 ✔
115	# `mod` is not a file so print an error message
116	print("Could not find the file called, `{}`\n".format(mod))	20 ✔
117	return None	20 ✔
118
119	# `mod` may be a relative path to a valid file so return its absolute path
120	return os.path.abspath(mod)	20 ✔
121
122
123	def _display(	20 ✔
124	cfgs: dict[nodes.NodeNG, ControlFlowGraph], filename: str, auto_open: bool = False
125	) -> None:
126	graph = graphviz.Digraph(name=filename + ".gv", **GRAPH_OPTIONS)	20 ✔
127	for node, cfg in cfgs.items():	20 ✔
128	if isinstance(node, nodes.Module):	20 ✔
129	subgraph_label = "__main__"	20 ✔
130	elif isinstance(node, nodes.FunctionDef):	20 ✔
131	scope_parent = node.scope().parent	20 ✔
132	subgraph_label = node.name	20 ✔
133	# Update the label to the qualified name if it is a method
134	if isinstance(scope_parent, nodes.ClassDef):	20 ✔
135	subgraph_label = scope_parent.name + "." + subgraph_label	20 ✔
136	else:
137	continue	×
138	with graph.subgraph(name=f"cluster_{cfg.cfg_id}") as c:	20 ✔
139	visited: set[str] = set()	20 ✔
140	_visit(cfg.start, c, visited, cfg.end)	20 ✔
141	for block in cfg.unreachable_blocks:	20 ✔
142	_visit(block, c, visited, cfg.end)	20 ✔
143	c.attr(label=subgraph_label, **SUBGRAPH_OPTIONS)	20 ✔
144
145	graph.render(outfile=filename + ".svg", view=auto_open)	20 ✔
146
147
148	def _visit(block: CFGBlock, graph: graphviz.Digraph, visited: set[str], end: CFGBlock) -> None:	20 ✔
149	"""
150	Visit a CFGBlock and add it to the control flow graph.
151	"""
152	node_id = f"{graph.name}_{block.id}"	20 ✔
153	if node_id in visited:	20 ✔
154	return	20 ✔
155
156	label = ""	20 ✔
157	fill_color = "white"	20 ✔
158
159	# Identify special cases
160	if len(block.statements) == 1:	20 ✔
161	stmt = block.statements[0]	20 ✔
162	if isinstance(stmt, nodes.Arguments):	20 ✔
163	label = f"{stmt.as_string()}\n"	20 ✔
164	fill_color = "palegreen"	20 ✔
165	elif isinstance(stmt.parent, nodes.If) and stmt is stmt.parent.test:	20 ✔
166	label = f"< if<U><B>{html.escape(stmt.as_string())}</B></U><BR/> >"	20 ✔
167	elif isinstance(stmt.parent, nodes.While) and stmt is stmt.parent.test:	20 ✔
168	label = f"< while<U><B>{html.escape(stmt.as_string())}</B></U><BR/> >"	20 ✔
169	elif isinstance(stmt.parent, nodes.For) and stmt is stmt.parent.iter:	20 ✔
170	label = f"< for {html.escape(stmt.parent.target.as_string())} in<U><B>{html.escape(stmt.as_string())}</B></U><BR/> >"	20 ✔
171	elif isinstance(stmt.parent, nodes.For) and stmt is stmt.parent.target:	20 ✔
172	label = f"< for<U><B>{html.escape(stmt.as_string())} </B></U> in {html.escape(stmt.parent.iter.as_string())}<BR/> >"	20 ✔
173
174	if block.statements and isinstance(block.statements[0], nodes.Pattern):	20 ✔
175	label = f"case {html.escape(block.statements[0].as_string())}"	×
176	label += f" if {block.statements[1].as_string()}" if len(block.statements) == 2 else ""	×
177
178	if not label: # Default	20 ✔
179	label = "\n".join([s.as_string() for s in block.statements]) + "\n"	20 ✔
180
181	# Need to escape backslashes explicitly.
182	label = label.replace("\\", "\\\\")	20 ✔
183	# \l is used for left alignment.
184	label = label.replace("\n", "\\l")	20 ✔
185
186	# Change the fill colour if block is the end of the cfg or unreachable
187	if block == end:	20 ✔
188	fill_color = "black"	20 ✔
189	elif not block.reachable:	20 ✔
190	fill_color = "grey93"	20 ✔
191
192	graph.node(node_id, label=label, fillcolor=fill_color, style="filled")	20 ✔
193	visited.add(node_id)	20 ✔
194
195	for edge in block.successors:	20 ✔
196	color = "black" if edge.is_feasible else "lightgrey"	20 ✔
197	if edge.get_label() is not None:	20 ✔
198	graph.edge(	20 ✔
199	node_id, f"{graph.name}_{edge.target.id}", label=edge.get_label(), color=color
200	)
201	else:
202	graph.edge(node_id, f"{graph.name}_{edge.target.id}", color=color)	20 ✔
203	_visit(edge.target, graph, visited, end)	20 ✔

pyta-uoft / pyta / 36646375051

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