How to match and modify a sequence of nodes ?

[Josverl]

I'm trying to find a solution to match and manipulate a combination of python commands/libcst.nodes from a file / module
this is rather trivial in regex, but i've not found a way to do the same using libcst

example:
from the code

x = 1

bar: int
"literal docstring"

I would like to identify/extract/manipulate the pair(s) of lines that contain an Assign or AnnAssign followed by a docstring

• bar: int
• "literal docstring"

import libcst as cst
import libcst.matchers as m

source = """\
x = 1

bar: int
"literal docstring"
"""
module = cst.parse_module(source)

# Matcher for (a module with)
# - an assignment or Annonated asignment of a literal
# - a docstring
literal_aa_matcher = m.Module( body= [
    m.ZeroOrMore(),
    m.SimpleStatementLine(body= [m.AnnAssign(),]),
    m.SimpleStatementLine(
        body= [m.Expr( value=m.SimpleString()),]
    ),
    m.ZeroOrMore(),
    ])

# try to match 
m.matches(module, literal_aa_matcher)
#>> match returns True 
m.findall(module, literal_aa_matcher)
#>> find returns the entire module , not usefull 
m.extract(module, literal_aa_matcher)
#>> returns {}
m.extractall(module, literal_aa_matcher)
#>> returns [{}]

I have tried many combinations including to try to match two (or more) consecutive nodes,
but I have not found a way to do this.

Matching directly on the body sequence or children of the module does not seem to work either.

Any suggestions appreceated

[zsol]

Maybe this is doable with the matcher API, but I'd personally just implement visit_SimpleStatementLine and have it keep track of the previous occurrence. This imperative approach seems simpler

[Josverl]

Thanks for the suggestion , didn't think of that as I was worried that the body structure was too complex for such an approach.
However as the definition is as simple as body: Sequence[Union[SimpleStatementLine, BaseCompoundStatement]] that should indeed be quite possible.

• visit_SimpleStatementLine
  if self.previous_node == AnnAssign or Assing and node ==SimpleString
  # save position / sequence
  else
  self.previous_node = node
• visit_BaseCompoundStatement
  self.previous_node = None

I think if I need to modify the first line however, that then I'd also need to add a metadata wrapper such as position to re-identify that same line on a second pass.

Option 2
another option I considered is to just iterate over the Module.body sequence,
and build a new body sequence from that.
with that body sequence instanciate an updated_module using that new body , and all other attributes of the original body.
That has the simplicity of accessing the nodes in the sequence by index [n-1] and does not require 2 visitor passes.

would that work, or am i overlooking something here ?

[zsol]

Yeah your option 2 sounds good too, probably better if you wanna change the first node.
Keep in mind that string node will be wrapped in at least one statement node in the body sequence