Comment python script programatically

My script file has lots of classes:

def func1():
  ...

class IPSEC(aetest.Testcase):
  ....

class IPSEC_DPI(aetest.Testcase):
  ....

class IPSEC_FNF(aetest.Testcase):
  ....
   .
   .
more classes
   .
   .
class all_results(aetest.Testcase):
  ....

all_results is the last class.

I need to create a function that accepts the class name and file path of the script above and comment that particular class. all_results will never be commented. I know its weird but thats the requirement as of now.

I created the function below, but it fails. Any other short method to achieve the same?

def comment_lines(file_name, profile):
    lines = open(file_name, 'r').readlines()
    start=None
    end=None
    i=0
    print(f"looking for class {profile}(aetest.Testcase):")
    for i in range(0, len(lines)):
        if f"class {profile}(aetest.Testcase):" in lines[i]:
            start=i
            break
        i+=1
    for i in range(start+1, len(lines)):
        if f"class " in lines[i]:
            end=i-2
            break
        i+=1
    for i in range(start, end+1):
        lines[i] = "#" + lines[i]
    out = open(file_name, 'w')
    out.writelines(lines)
    out.close()

Do It The Right Way™ with libcst and the structural pattern matching feature introduced in Python 3.10. If updating to 3.10 is not a choice, if/elif/else blocks can be used instead, but that is an exercise for the reader.

<pre><code>import libcst as cst

Naming is hard

class ClassFile:
def init(self, file_name: str):
self._file_name = file_name

def comment_out(self, class_name):
with open(self._file_name, 'r+') as file:
code = file.read()
tree = cst.parse_module(code)

  # <a href="https://libcst.readthedocs.io/en/latest/nodes.html#libcst.Module.visit">.visit()</a> calls the visitor's methods
# every time it visits and leaves a node.
new_tree = tree.visit(ClassDefCommenter(tree, class_name))
file.seek(0)
file.write(new_tree.code)

</code></pre>

The visitor we need is an instance of a subclass of cst.CSTTransformer:

class ClassDefCommenter(cst.CSTTransformer):

  def __init__(self, module: cst.Module, class_name: str):
    self._module = module
    self._class_name = class_name
  
  def leave_ClassDef(self, node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.CSTNode:
    '''
    This method takes in two argument:
    - The original node, and
    - a deep copy of that node (to be inserted to the new tree).

    We need to find a class definition block whose:
    - .name is the given class name
    - .bases is a sequence with only one argument: 'aetest.Testcase'
    '''

    match node:
      case cst.ClassDef(
        name = cst.Name(value = self._class_name),
        bases = [
          cst.Arg(
            value = cst.Attribute(
              value = cst.Name(value = 'aetest'),
              attr = cst.Name(value = 'Testcase')
            )
          )
        ]
      ):
        # Replace any such block with a multiline string
        # containing the original source code.
        return cst.SimpleString(
          value = f"'''{self._module.code_for_node(node)}'''"
        )
      case _:
        # Else, do nothing.
        return updated_node

For simplicity and ease of explanation, I used multiline string syntax instead of comment. You might want to escape nested triple quotes, if any.

Try it:

ClassFile('filename.py').comment_out('IPSEC')

This file:

class aetest:

  class Testcase:
    ...

def func1():
  ...

class IPSEC(aetest.Testcase):
  ...
	
class IPSEC_DPI(aetest.Testcase):
  ...

class IPSEC_FNF(aetest.Testcase):
  ...

class all_results(aetest.Testcase):
  ...

...will be modified to:

class aetest:

  class Testcase:
    ...

def func1():
  ...
'''
class IPSEC(aetest.Testcase):
  ...
'''	
class IPSEC_DPI(aetest.Testcase):
  ...

class IPSEC_FNF(aetest.Testcase):
  ...

class all_results(aetest.Testcase):
  ...

How about this?

def comment_lines(file_name, profile):
    with open(file_name, 'r') as f:
        lines = f.readlines()

    profile = f"class {profile}(aetest.Testcase):"
    comment = False
    print(f"looking for {profile}")
    for i in range(len(lines)):
        if  lines[i].strip() == profile in lines[i]:
             comment = True
        elif comment and lines[i].startswith("class"):
             comment = False
    
        if comment:
             lines[i] = '# ' + lines[i]
    
    with open(file_name, 'w') as f:
        f.writelines(lines)

Some tips that I should mention:

1. You don't need to define the variable you want to use in a for-loop before like other languages (C, C++, Java, etc) Python will do it for you!
2. Use with to open a file and read/write in them as it will close and save them automatically.
3. Try doing everything in one loop as it is much faster!
4. Use or update variables like profile before loop to make the loop faster!

Another implementation using a tri-state flag found:

def comment_lines(file_name, profile):
with open(file_name) as f:
lines = f.readlines()
with open(file_name, 'w') as f:
found = 0
for line in lines:
if found == 0 and f"class {profile}(aetest.Testcase):" in line:
line = '#' + line
found = 1
elif found == 1:
if 'class' in line:
found = 2
else:
line = '#' + line
f.write(line)

And another implementation, slightly longer, but that doesn't comment blank lines after the end of the target class definition (it still comments blank lines inside said definition); this is achieved by bufferizing blank lines and writing them to the file only after a not-blank line has been encountered:

def comment_lines(file_name, profile):
with open(file_name) as f:
lines = f.readlines()
with open(file_name, 'w') as f:
found = False
buffer = []
for line in lines:
if f"class {profile}(aetest.Testcase):" in line:
found = True
elif found:
if not line.strip():
buffer.append(line)
continue
elif 'class' in line:
found = False
for l in buffer:
f.write('# '*found + l)
buffer = []
f.write('# '*found + line)