在文本文件指定位置插入行

Question

我有一个文本文件，内容大概是这样的：

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

现在我想在 'foo1 bar3' 和 'foo2 bar4' 之间插入 'foo bar'。

我这样做的：

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)

这个方法对我有效，但看起来有点丑。

Answer 1

记住，迭代器是一种重要的对象。它可以在多个for循环中使用。

这里有一种方法，可以在不使用很多复杂的if语句和标志的情况下来处理这个问题。

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()

Answer 2

这是在调整Alex Martelli的例子：

import fileinput
for line in fileinput.input('1.txt', inplace=1):
 print line,
 if line.startswith('foo1 bar3'):
     print 'foo bar'

Answer 3

在Python中，想要对文件进行“伪就地”修改，最好的方法是使用标准库中的fileinput模块：

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,

如果你想保留旧版本，可以指定一个备份扩展名，这样就能保存之前的文件。这个方法和你的代码类似——使用.bak作为备份扩展名，但在修改成功后会自动删除这个备份。

除了使用正确的标准库模块，这段代码的逻辑也更简单：它的目的是在每一行以foo1开头的行后面插入一行"foo bar"。只需要一个布尔值（我现在是否在这样的行里面？）就可以了。这个布尔值可以根据当前行是否以foo1开头来简单设置。如果你想要的具体逻辑和这个稍有不同（这是我从你的代码中推测的），调整这段代码应该也不难。