[question] Exclude If in `elif` position

[Feuermurmel]

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:

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:

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:

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

Is there any way to do this?

[frvnkliu]

Could you just make a helper function that calls the additional statement and returns the attribute then replace all accesses of the attribute with this new helper function?

[Feuermurmel]

Could you just make a helper function that calls the additional statement and returns the attribute then replace all accesses of the attribute with this new helper function?

I could, probably. I do not remember what code exactly I wanted to insert, but I believe it was something that interacted with type checking, like reveal_type(), assert isinstance(...), or assert_type(). Those could all be extracted into a function or placed inline, but that would make the diff harder to read and harder to revert. The inserted code should only be temporary.