CODEWITHSUNDEEP
pythonscripting
InvertReality
InvertReality

Reputation: 39

Python Increasing counter by 1 every time something is written to a file

What I'm trying to do is every time the file 'TCPStreamPayload' is written to I want a counter that increases every time that information is appended to the file.

For example I have:

payloadRecon = open('TCPStreamPayload.txt','a')

segmentCount = 0

--my script here--

then at the bottom

if payloadRecon.write:
    segmentCount +=1
payloadRecon.write('####TCP PAYLOAD STREAM RECONSTRUCTION. Segment ' + str(segmentCount) + '######' '\n\n\n\n').

Basically, everytime I write to this file using my script I want the segment number to increase by 1.

Upvotes: 2

Views: 10742

Answers (1)

abarnert
abarnert

Reputation: 365707

You really should write a wrapper class and use that instead of a plain file:

class CountingWrapper(object):
    def __init__(self, *args, **kwargs):
        self.segment_count = 0
        self.f = open(*args, **kwargs)
    def write(self, *args, **kwargs):
        self.segment_count += 1
        return self.f.write(*args, **kwargs)

payloadRecon = CountingWrapper('TCPStreamPayload.txt', 'a')
payloadRecon.write('dsfsd')
assert payloadRecon.count == 1

Alternatively, if you're just trying to wrap up the three lines of code so you don't have to keep repeating yourself, you can just write a wrapper function:

segmentCount = 0
def payload_write(f, *args, **kwargs):
    global segmentCount
    segmentCount += 1
    return f.write(*args, **kwargs)

Now, you just need this one-liner to count and write:

payload_write(payloadRecon, '####TCP PAYLOAD STREAM RECONSTRUCTION. Segment ' + str(segmentCount) + '######' '\n\n\n\n')
assert segmentCount == 1

The class is more flexible, if you need to wrap multiple methods. But you obviously need to delegate any other methods you care about. Of course you could do the delegation dynamically, via __getattr__. Or you could use inheritance. (But keep in mind that in 3.x, the class is not file, but one of the classes in the io package, and you're not supposed to create either file or io.whatever by constructor anyway, but by open, so you'll still need to write a wrapper factory function to replace open. The good news is, that means you don't need to handle __new__, just __init__.)

In 3.x, you could even monkeypatch the instance (but this won't work in 2.x, where file.write is read-only):

payloadRecon = open('TCPStreamPayload.txt', 'a')
segmentCount = 0
write_original = payloadRecon.write
def write_and_count(self, *args, **kwargs):
    global segmentCount
    segmentCount += 1
    write_original(self, *args, **kwargs)
payloadRecon.write = write_and_count

But really, keeping things explicit is probably the way to go here.

Upvotes: 5

Related Questions