Inserting Line at Specified Position of a Text File in Python

Question

Hi,

I have a text file which looks like this:

blah blah
foo1 bar1
foo1 bar2
foo1 bar3
foo2 bar4
foo2 bar5
blah blah

Now I want to insert 'foo bar' between 'foo1 bar3' and 'foo2 bar4'.

This is how I did it:

import shutil

txt = '1.txt'
tmptxt = '1.txt.tmp'

with open(tmptxt, 'w') as outfile:
    with open(txt, 'r') as infile:
        flag = 0
        for line in infile:
            if not line.startswith('foo1') and flag == 0:
                outfile.write(line)
                continue
            if line.startswith('foo1') and flag == 0:
                flag = 1
                outfile.write(line)
                continue
            if line.startswith('foo1') and flag == 1:
                outfile.write(line)
                continue
            if not line.startswith('foo1') and flag == 1:
                outfile.write('foo bar\n')
                outfile.write(line)
                flag = 2
                continue
            if not line.startswith('foo1') and flag == 2:
                outfile.write(line)
                continue

shutil.move(tmptxt, txt)

This works for me, but looks rather ugly. Please help me improve it, or throw me some ideas so I can write it again.

Many thanks.

Answer 1

The best way to make "pseudo-inplace" changes to a file in Python is with the fileinput module from the standard library:

import fileinput

processing_foo1s = False

for line in fileinput.input('1.txt', inplace=1):
  if line.startswith('foo1'):
    processing_foo1s = True
  else:
    if processing_foo1s:
      print 'foo bar'
    processing_foo1s = False
  print line,

You can also specify a backup extension if you want to keep the old version around, but this works in the same vein as your code -- uses .bak as the backup extension but also removes it once the change has successfully completed.

Besides using the right standard library module, this code uses simpler logic: to insert a "foo bar" line after every run of lines starting with foo1, a boolean is all you need (am I inside such a run or not?) and the bool in question can be set unconditionally just based on whether the current line starts that way or not. If the precise logic you desire is slightly different from this one (which is what I deduced from your code), it shouldn't be hard to tweak this code accordingly.

Answer 2

Recall that an iterator is a first-class object. It can be used in multiple for statements.

Here's a way to handle this without a lot of complex-looking if-statements and flags.

with open(tmptxt, 'w') as outfile:
    with open(txt, 'r') as infile:
        rowIter= iter(infile)
        for row in rowIter:
            if row.startswith('foo2'): # Start of next section
                 break
            print row.rstrip(), repr(row)
        print "foo bar"
        print row
        for row in rowIter:
            print row.rstrip()