I want to add a comment mark to the end of every if statement in python code, how can I do it?

Question

Here's a piece of code that doesn't make sense:

if aaa=1:
# dfdf fdf
  xxx
  if cc = 4:
    rrr
  else:
    ww
    if aaaaa:
      # dssfdsf
      ttt
    # endif
  # endif
elif bb=2:

  if ll is True:
    return
  # endif
  yyy
else:
  zzz
# endif
dfsdfsaa

The #endif is the comment mark I want to add, I tried to do it with regular expressions ( regex101Example ), but it didn't work, who has a better idea?

Answer 1

The code below seems to be a solution:

# coding=utf8
# the above tag defines encoding for this document and is for Python 2.x compatibility

import re
test_str = ("if aaa=1:\n"
    "# dfdf fdf\n"
    "  xxx\n"
    "  if cc = 4:\n"
    "    rrr\n"
    "  else:\n"
    "    ww\n"
    "    if aaaaa:\n"
    "      # dssfdsf\n"
    "      ttt\n"
    "elif bb=2:\n\n"
    "  if ll is True:\n"
    "    return\n"
    "  yyy\n"
    "else:\n"
    "  zzz\n"
    "dfsdfsaa")
regex = r"^( *)(if(.|\n)+?^(?!\1  |#|\1el|\n))"
subst = "\\1#\\2\\1# endif\\n"
matches = re.match(regex, test_str, re.MULTILINE)
# Assuming a maximum nesting of 6 levels
for i in range(6):
  if matches:
    test_str = re.sub(regex, subst, test_str, 0, re.MULTILINE)
result = re.sub(r"#if ", "if ",test_str, )
print(result)

Modified according to Code Generator by regex101Example . I have been trying to implement directly through regex's recursion and other methods, but no success...

Output:

if aaa=1:
# dfdf fdf
  xxx
  if cc = 4:
    rrr
  else:
    ww
    if aaaaa:
      # dssfdsf
      ttt
    # endif
  # endif
elif bb=2:

  if ll is True:
    return
  # endif
  yyy
else:
  zzz
# endif
dfsdfsaa