Instagram / Instagram/LibCST

[question] Exclude If in `elif` position

Open
#1,286 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.9k
Forks
229
PR merge metrics
No merged PRs in 30d

Description

I'm trying to write a transformer that inserts a certain statement before existing statements that contains some code pattern, but I'm hitting a problem with `If` nodes that represent the `elif` part of an `if` statement.

More specifically, I'm trying to insert an additional statement whenever an attribute with a specific name on any object is accessed. The attribute access could be anywhere, inside a simple expression statement, the right side of an assignment, the condition of an `if`, etc.

Below is a simplified version of the code I'm using. It works well in almost all cases. Shown is the case of the attribute access happening in the condition of an `if` statement:

```python
from textwrap import dedent
from libcst import Attribute, CSTNode, BaseStatement, CSTTransformer, \
FlattenSentinel, parse_module, parse_statement
import libcst.matchers as m

class Transformer(CSTTransformer):
def __init__(self):
super().__init__()

# Attribute accesses within each nested BaseStatement node.
self.attr_accesses_stack: list[list[Attribute]] = []

def on_visit(self, node: CSTNode) -> bool:
if isinstance(node, BaseStatement):
self.attr_accesses_stack.append([])

if m.matches(node, m.Attribute(attr=m.Name("special_attribute"))):
self.attr_accesses_stack[-1].append(node)

return True

def on_leave(
self, original_node: CSTNode, updated_node: CSTNode
) -> CSTNode | FlattenSentinel[CSTNode]:
if isinstance(updated_node, BaseStatement):
attr_accesses = self.attr_accesses_stack.pop()
new_nodes = [updated_node]

# The actual code does something more complex with the collected attribute nodes.
for i in attr_accesses:
new_nodes.insert(0, parse_statement("print('attribute accessed')"))

return FlattenSentinel(new_nodes)

return updated_node

example_code = dedent(
"""\
if x.special_attribute:
pass
# elif y.special_attribute:
# pass
"""
)

print(parse_module(example_code).visit(Transformer()).code)
```

Output:

```python
print('attribute accessed')
if x.special_attribute:
pass
# elif y.special_attribute:
# pass
```

The problem arises when the attribute access is in the `elif`'s condition. Uncommenting the two lines in `example_code` produces the following error:

```
File [...]/venv/lib/python3.12/site-packages/libcst/_nodes/internal.py:112 in visit_optional
raise TypeError(

TypeError: We got a FlattenSentinel while visiting a If. This node's parent does not allow for it to be it to be replaced with a sequence.
```

The problem is clear: An `If` node is used in two distinct cases:

- As an item in a list of statements to represent the a whole `if` statement including any `elif` and `else` parts.
- To represent an `elif` part of an `if` statement. There can only be an `If`, an `Else` or no node in that place.

The second case is where my code fails because it returns a `FlattenSentinel` instance.

**My question:** What is the best way to handle this case? What I'd like to do is to ignore these `If` nodes in the traversal so that the attribute accesses are instead collected for the top-level `If` node of each `if` statement, i.e. so that the resulting code would look like this:

```python
print('attribute accessed')
print('attribute accessed')
if x.special_attribute:
pass
elif y.special_attribute:
pass
```

Is there any way to do this?

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.