206. AST (Abstract Syntax Trees) in Python

🔹 1. Parsing Python Code into an AST

• Convert a Python expression into an AST tree.

Copy

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.

Copy

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.

Copy

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.

Copy

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.

Copy

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.

Copy

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.

Copy

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.

Copy

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.

Copy

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.

Copy

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

SnippetDescription

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.