self modifying python script

Viewed 11294

I want to create python script which can modify code in that script itself using Python Language Services or using any other way.
e.g. A script which keep track of its count of successfull execution

import re
COUNT = 0

def updateCount():
    # code to update second line e.g. COUNT = 0
    pass

if __name__ == '__main__':
    print('This script has run {} times'.format(COUNT))
    updateCount()

On successful execution of this script code should get changed to

import re
COUNT = 1

def updateCount():
    # code to update second line e.g. COUNT = 0
    pass

if __name__ == '__main__':
    print('This script has run {} times'.format(COUNT))
    updateCount()

Simple approach came to my mind was to open __file__ in write mode and do requried modification using reguler expessions etc. But that did not work I got exception io.UnsupportedOperation: not readable. Even if this approach would be working then it would be very risky because it can spoil my whole script. so I am looking for solution using Python Language Services.

4 Answers

@kyriakosSt's answer works but hard-codes that the assignment to COUNT must be on the second line, which can be prone to unexpected behaviors over time when the line number changes due to the source being modified for something else.

For a more robust solution, you can use lib2to3 to parse and update the source code instead, by subclassing lib2to3.refactor.RefactoringTool to refactor the code using a fixer that is a subclass of lib2to3.fixer_base.BaseFix with a pattern that looks for an expression statement with the pattern 'COUNT' '=' any, and a transform method that updates the last child node by incrementing its integer value:

from lib2to3 import fixer_base, refactor

COUNT = 0 # this should be incremented every time the script runs

class IncrementCount(fixer_base.BaseFix):
    PATTERN = "expr_stmt< 'COUNT' '=' any >"

    def transform(self, node, results):
        node.children[-1].value = str(int(node.children[-1].value) + 1)
        return node

class Refactor(refactor.RefactoringTool):
    def __init__(self, fixers):
        self._fixers = [cls(None, None) for cls in fixers]
        super().__init__(None)

    def get_fixers(self):
        return self._fixers, []

with open(__file__, 'r+') as file:
    source = str(Refactor([IncrementCount]).refactor_string(file.read(), ''))
    file.seek(0)
    file.write(source)

Demo: https://repl.it/@blhsing/MushyStrangeClosedsource

This will edit the module level variables defined before _local_config. Later, process an update to the dictionary, then replace the line when iterating over the source file with the new _local_config values:

count = 0
a = 0
b = 1
c = 1

_local_config = dict(
    filter(
        lambda elem: (elem[0][:2] != "__") and (str(elem[1])[:1] != "<"),
        globals().items(),
    ),
)

# do some stuff
count += 1
c = a + b
a = b
b = c

# update with new values
_local_config = dict(
    filter(
        lambda elem: elem[0] in _local_config.keys(),
        globals().items(),
    )
)

# read self
with open(__file__, "r") as f:
    new_file = ""

    for line in f.read().split("\n"):
        for k, v in _local_config.items():

            search = f"{k} = "

            if search == line[: len(k) + 3]:
                line = search + str(v)
                _local_config.pop(k)
                break

        new_file += line + "\n"

# write self
with open(__file__, "w") as f:
    f.write(new_file[:-1])
Related