Skip to content

Commit 91f42be

Browse files
hunterhoganhunterhogan
committed
typing_extensions.TypeIs
1 parent e78ba21 commit 91f42be

1 file changed

Lines changed: 132 additions & 120 deletions

File tree

‎astToolkit/_toolkitNodeVisitor.py‎

Lines changed: 132 additions & 120 deletions
Original file line numberDiff line numberDiff line change
@@ -27,132 +27,144 @@
2727
seamlessly with astToolkit's type system and atomic classes to create composable, maintainable AST manipulation
2828
code.
2929
"""
30+
3031
from astToolkit import 个return, 木
3132
from collections.abc import Callable
32-
from typing import Any, cast, Generic, TypeIs
33+
from typing import Any, cast, Generic
34+
from typing_extensions import TypeIs
3335
import ast
3436

3537
# TODO Identify the logic that narrows the type and can help the user during static type checking.
3638

39+
3740
class NodeTourist(ast.NodeVisitor, Generic[木, 个return]):
38-
"""
39-
Read-only AST visitor that extracts information from nodes matching predicate conditions.
40-
(AI generated docstring)
41-
42-
`NodeTourist` implements the antecedent-action pattern for non-destructive AST analysis. It traverses an AST
43-
tree, applies predicate functions to identify target nodes, and executes action functions on matches to extract
44-
or analyze information. The visitor preserves the original AST structure while capturing results from matching
45-
nodes.
46-
47-
This class is particularly useful for analysis workflows where you need to gather information about specific
48-
node types or patterns without modifying the source code structure. The generic type parameters ensure type
49-
safety when working with specific AST node types and return values.
50-
51-
Parameters:
52-
findThis: Predicate function that tests AST nodes. Can return either a `TypeIs` for type narrowing or a
53-
simple boolean. When using `TypeIs`, the type checker can safely narrow the node type for the action
54-
function.
55-
doThat: Action function that operates on nodes matching the predicate. Receives the matched node with
56-
properly narrowed typing and returns the extracted information.
57-
58-
Examples:
59-
Extract all function names from a module:
60-
```python
61-
functionNameCollector = NodeTourist(Be.FunctionDef, lambda functionDef: DOT.name(functionDef))
62-
functionNames = []
63-
functionNameCollector.doThat = Then.appendTo(functionNames)
64-
functionNameCollector.visit(astModule)
65-
```
66-
67-
Find specific function definition:
68-
```python
69-
specificFunction = NodeTourist(IfThis.isFunctionDefIdentifier("targetFunction"), Then.extractIt)
70-
foundFunction = specificFunction.captureLastMatch(astModule)
71-
```
72-
"""
73-
def __init__(self, findThis: Callable[[ast.AST], TypeIs[木] | bool], doThat: Callable[[木], 个return]) -> None:
74-
self.findThis = findThis
75-
self.doThat = doThat
76-
self.nodeCaptured: 个return | None = None
77-
78-
def visit(self, node: ast.AST) -> None:
79-
if self.findThis(node):
80-
node = cast(木, node)
81-
self.nodeCaptured = self.doThat(node)
82-
self.generic_visit(node)
83-
84-
def captureLastMatch(self, node: ast.AST) -> 个return | None:
85-
"""
86-
Visit an AST tree and return the result from the last matching node.
87-
(AI generated docstring)
88-
89-
This method provides a convenient interface for single-result extraction workflows. It resets the internal
90-
capture state, traverses the provided AST tree, and returns the result from the most recently matched node.
91-
If no nodes match the predicate, returns `None`.
92-
93-
The method is particularly useful when you expect exactly one match or when you only care about the final
94-
match in traversal order. For collecting multiple matches, modify the `doThat` action function to append
95-
results to a collection.
96-
97-
Parameters:
98-
node: Root AST node to begin traversal from. Can be any AST node type including modules, functions,
99-
classes, or expressions.
100-
101-
Returns:
102-
lastResult: Result from the action function applied to the last matching node, or `None` if no matches
103-
were found during traversal.
104-
"""
105-
self.nodeCaptured = None
106-
self.visit(node)
107-
return self.nodeCaptured
41+
"""
42+
Read-only AST visitor that extracts information from nodes matching predicate conditions.
43+
(AI generated docstring)
44+
45+
`NodeTourist` implements the antecedent-action pattern for non-destructive AST analysis. It traverses an AST
46+
tree, applies predicate functions to identify target nodes, and executes action functions on matches to extract
47+
or analyze information. The visitor preserves the original AST structure while capturing results from matching
48+
nodes.
49+
50+
This class is particularly useful for analysis workflows where you need to gather information about specific
51+
node types or patterns without modifying the source code structure. The generic type parameters ensure type
52+
safety when working with specific AST node types and return values.
53+
54+
Parameters:
55+
findThis: Predicate function that tests AST nodes. Can return either a `TypeIs` for type narrowing or a
56+
simple boolean. When using `TypeIs`, the type checker can safely narrow the node type for the action
57+
function.
58+
doThat: Action function that operates on nodes matching the predicate. Receives the matched node with
59+
properly narrowed typing and returns the extracted information.
60+
61+
Examples:
62+
Extract all function names from a module:
63+
```python
64+
functionNameCollector = NodeTourist(Be.FunctionDef, lambda functionDef: DOT.name(functionDef))
65+
functionNames = []
66+
functionNameCollector.doThat = Then.appendTo(functionNames)
67+
functionNameCollector.visit(astModule)
68+
```
69+
70+
Find specific function definition:
71+
```python
72+
specificFunction = NodeTourist(IfThis.isFunctionDefIdentifier("targetFunction"), Then.extractIt)
73+
foundFunction = specificFunction.captureLastMatch(astModule)
74+
```
75+
"""
76+
77+
def __init__(
78+
self,
79+
findThis: Callable[[ast.AST], TypeIs[木] | bool],
80+
doThat: Callable[[木], 个return],
81+
) -> None:
82+
self.findThis = findThis
83+
self.doThat = doThat
84+
self.nodeCaptured: 个return | None = None
85+
86+
def visit(self, node: ast.AST) -> None:
87+
if self.findThis(node):
88+
node = cast(木, node)
89+
self.nodeCaptured = self.doThat(node)
90+
self.generic_visit(node)
91+
92+
def captureLastMatch(self, node: ast.AST) -> 个return | None:
93+
"""
94+
Visit an AST tree and return the result from the last matching node.
95+
(AI generated docstring)
96+
97+
This method provides a convenient interface for single-result extraction workflows. It resets the internal
98+
capture state, traverses the provided AST tree, and returns the result from the most recently matched node.
99+
If no nodes match the predicate, returns `None`.
100+
101+
The method is particularly useful when you expect exactly one match or when you only care about the final
102+
match in traversal order. For collecting multiple matches, modify the `doThat` action function to append
103+
results to a collection.
104+
105+
Parameters:
106+
node: Root AST node to begin traversal from. Can be any AST node type including modules, functions,
107+
classes, or expressions.
108+
109+
Returns:
110+
lastResult: Result from the action function applied to the last matching node, or `None` if no matches
111+
were found during traversal.
112+
"""
113+
self.nodeCaptured = None
114+
self.visit(node)
115+
return self.nodeCaptured
116+
108117

109118
class NodeChanger(ast.NodeTransformer):
110-
"""
111-
Destructive AST transformer that selectively modifies nodes matching predicate conditions.
112-
(AI generated docstring)
113-
114-
`NodeChanger` implements the antecedent-action pattern for targeted AST transformation. It extends Python's
115-
`ast.NodeTransformer` to provide precise control over which nodes are modified during tree traversal. The
116-
transformer applies predicate functions to identify target nodes and executes action functions to perform
117-
modifications, replacements, or deletions.
118-
119-
This class forms the foundation for code optimization, refactoring, and generation workflows. Unlike
120-
`NodeTourist`, `NodeChanger` modifies the AST structure and returns a transformed tree. The transformation
121-
is applied recursively, ensuring that nested structures are properly processed.
122-
123-
The class is designed for scenarios where you need to make surgical changes to specific parts of an AST while
124-
preserving the overall structure and semantics of the code. Common use cases include function inlining,
125-
variable renaming, dead code elimination, and pattern-based code transformations.
126-
127-
Parameters:
128-
findThis: Predicate function that identifies nodes to transform. Should return `True` for nodes that require
129-
modification and `False` for nodes that should remain unchanged. The function receives an `ast.AST` node
130-
and determines whether transformation is needed.
131-
doThat: Action function that performs the actual transformation. Receives nodes that matched the predicate
132-
and returns the replacement node, modified node, or `None` for deletion. The return value becomes the new
133-
node in the transformed tree.
134-
135-
Examples:
136-
Replace all function calls to a specific function:
137-
```python
138-
callReplacer = NodeChanger(
139-
IfThis.isCallIdentifier("oldFunction"),
140-
Then.replaceWith(Make.Call(Make.Name("newFunction"), [], []))
141-
)
142-
transformedAST = callReplacer.visit(originalAST)
143-
```
144-
145-
Remove all pass statements:
146-
```python
147-
passRemover = NodeChanger(Be.Pass, Then.removeIt)
148-
cleanedAST = passRemover.visit(originalAST)
149-
```
150-
"""
151-
def __init__(self, findThis: Callable[..., Any], doThat: Callable[..., Any]) -> None:
152-
self.findThis = findThis
153-
self.doThat = doThat
154-
155-
def visit(self, node: ast.AST) -> ast.AST:
156-
if self.findThis(node):
157-
return self.doThat(node)
158-
return super().visit(node)
119+
"""
120+
Destructive AST transformer that selectively modifies nodes matching predicate conditions.
121+
(AI generated docstring)
122+
123+
`NodeChanger` implements the antecedent-action pattern for targeted AST transformation. It extends Python's
124+
`ast.NodeTransformer` to provide precise control over which nodes are modified during tree traversal. The
125+
transformer applies predicate functions to identify target nodes and executes action functions to perform
126+
modifications, replacements, or deletions.
127+
128+
This class forms the foundation for code optimization, refactoring, and generation workflows. Unlike
129+
`NodeTourist`, `NodeChanger` modifies the AST structure and returns a transformed tree. The transformation
130+
is applied recursively, ensuring that nested structures are properly processed.
131+
132+
The class is designed for scenarios where you need to make surgical changes to specific parts of an AST while
133+
preserving the overall structure and semantics of the code. Common use cases include function inlining,
134+
variable renaming, dead code elimination, and pattern-based code transformations.
135+
136+
Parameters:
137+
findThis: Predicate function that identifies nodes to transform. Should return `True` for nodes that require
138+
modification and `False` for nodes that should remain unchanged. The function receives an `ast.AST` node
139+
and determines whether transformation is needed.
140+
doThat: Action function that performs the actual transformation. Receives nodes that matched the predicate
141+
and returns the replacement node, modified node, or `None` for deletion. The return value becomes the new
142+
node in the transformed tree.
143+
144+
Examples:
145+
Replace all function calls to a specific function:
146+
```python
147+
callReplacer = NodeChanger(
148+
IfThis.isCallIdentifier("oldFunction"),
149+
Then.replaceWith(Make.Call(Make.Name("newFunction"), [], []))
150+
)
151+
transformedAST = callReplacer.visit(originalAST)
152+
```
153+
154+
Remove all pass statements:
155+
```python
156+
passRemover = NodeChanger(Be.Pass, Then.removeIt)
157+
cleanedAST = passRemover.visit(originalAST)
158+
```
159+
"""
160+
161+
def __init__(
162+
self, findThis: Callable[..., Any], doThat: Callable[..., Any]
163+
) -> None:
164+
self.findThis = findThis
165+
self.doThat = doThat
166+
167+
def visit(self, node: ast.AST) -> ast.AST:
168+
if self.findThis(node):
169+
return self.doThat(node)
170+
return super().visit(node)

0 commit comments

Comments
 (0)