206. AST (Abstract Syntax Trees) in Python

🔹 1. Parsing Python Code into an AST

• Convert a Python expression into an AST tree.

import ast

code = "x = 5 + 3"
tree = ast.parse(code)

print(ast.dump(tree, indent=4))

🔍 How it works:

• ast.parse(code) generates an AST representation of the code.

• ast.dump(tree, indent=4) prints the AST structure in a readable format.

---

🔹 2. Extracting Function Names from Python Code

• Traverse an AST to find all function definitions.

import ast

code = """
def greet():
    print("Hello, World!")

def add(a, b):
    return a + b
"""

tree = ast.parse(code)

for node in ast.walk(tree):
    if isinstance(node, ast.FunctionDef):
        print("Function name:", node.name)

🔍 How it works:

• Uses ast.walk(tree) to iterate over all nodes.

• Checks for ast.FunctionDef nodes to find function definitions.

---

🔹 3. Counting Number of If-Statements in Code

• Traverse the AST to count if statements.

import ast

code = """
if x > 10:
    print("Greater")
elif x == 10:
    print("Equal")
else:
    print("Smaller")
"""

tree = ast.parse(code)

count = sum(1 for node in ast.walk(tree) if isinstance(node, ast.If))
print("Number of if-statements:", count)

🔍 How it works:

• Uses sum(1 for node in ast.walk(tree) if isinstance(node, ast.If)) to count if nodes.

---

🔹 4. Extracting All Variable Assignments

• Get all variables assigned in the code.

import ast

code = """
x = 10
y = x + 5
name = "Alice"
"""

tree = ast.parse(code)

variables = [node.targets[0].id for node in ast.walk(tree) if isinstance(node, ast.Assign)]
print("Assigned variables:", variables)

🔍 How it works:

• Looks for ast.Assign nodes and extracts variable names.

---

🔹 5. Detecting Function Calls in Code

• Identify all function calls in a script.

import ast

code = """
print("Hello")
result = max(3, 5)
"""

tree = ast.parse(code)

calls = [node.func.id for node in ast.walk(tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)]
print("Function calls:", calls)

🔍 How it works:

• Detects ast.Call nodes and extracts function names.

---

🔹 6. Modifying AST – Changing Function Names

• Rewrite function names dynamically.

import ast

class RenameFunctions(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        node.name = "modified_" + node.name
        return node

code = """
def hello():
    return "Hi"

def add(a, b):
    return a + b
"""

tree = ast.parse(code)
tree = RenameFunctions().visit(tree)

print(ast.dump(tree, indent=4))  # Modified AST

🔍 How it works:

• Uses ast.NodeTransformer to rename all function definitions.

---

🔹 7. Executing Modified AST Code

• Convert modified AST back to executable code.

import ast

code = """
def greet():
    return "Hello"
"""

tree = ast.parse(code)

# Convert AST back to source code
modified_code = ast.unparse(tree)
exec(modified_code)  # Runs the code
print(greet())  # Output: Hello

🔍 How it works:

• ast.unparse(tree) converts AST back to code.

• exec() executes the modified code.

---

🔹 8. Checking for Unsafe exec() Usage

• Detects the use of exec() in user-supplied code.

import ast

code = """
exec('print("Malicious Code")')
"""

tree = ast.parse(code)

if any(isinstance(node, ast.Exec) for node in ast.walk(tree)):
    print("Unsafe exec() detected!")
else:
    print("Code is safe.")

🔍 How it works:

• Scans for ast.Exec nodes to detect exec() calls.

---

🔹 9. Extracting All Imported Modules

• Find all modules imported in a script.

import ast

code = """
import os
import sys
from math import sqrt
"""

tree = ast.parse(code)

imports = [node.names[0].name for node in ast.walk(tree) if isinstance(node, ast.Import)]
from_imports = [node.module for node in ast.walk(tree) if isinstance(node, ast.ImportFrom)]

print("Imports:", imports + from_imports)

🔍 How it works:

• Extracts all import and from ... import statements.

---

🔹 10. Detecting List Comprehensions

• Identify list comprehensions in the code.

import ast

code = """
nums = [x for x in range(10)]
squares = [x**2 for x in nums]
"""

tree = ast.parse(code)

count = sum(1 for node in ast.walk(tree) if isinstance(node, ast.ListComp))
print("Number of list comprehensions:", count)

🔍 How it works:

• Counts occurrences of ast.ListComp, which represents list comprehensions.

---

Summary of AST Python Code Snippets

Snippet	Description
1️⃣	Parsing Python code into an AST
2️⃣	Extracting function names from AST
3️⃣	Counting if statements in code
4️⃣	Finding assigned variables
5️⃣	Detecting function calls
6️⃣	Modifying function names dynamically
7️⃣	Executing modified AST code
8️⃣	Checking for unsafe exec() usage
9️⃣	Extracting all imported modules
🔟	Detecting list comprehensions

---

🚀 Conclusion

• AST (Abstract Syntax Trees) allow deep code analysis and transformations.

• Python’s ast module provides parsing, modification, and execution of AST objects.

• This is useful for security checks, refactoring, and optimizations.