generator to write N lines to file

[Peter Otten]

Sayth Renshaw wrote:

> Afternoon
> 
> Trying to create a generator to write the first N lines of text to a file.
> However, I keep receiving the islice object not the text, why?
> 
> from itertools import islice
> 
> fileName = dir / "7oldsamr.txt"
> dumpName = dir / "dump.txt"
> 
> def getWord(infile, lineNum):
>     with open(infile, 'r+') as f:
>         lines_gen = islice(f, lineNum)
>         yield lines_gen

To get the individual lines you have to yield them

          for line in lines_gen:
              yield line

This can be rewritten with some syntactic sugar as

          yield from lines_gen
 
> for line in getWord(fileName, 5):
>     with open(dumpName, 'a') as f:
>         f.write(line)

Note that getWord as written does not close the infile when you break out of 
the loop, e. g.

for line in getWord(filename, 5):
    break

To avoid that you can use a context manager:

from itertools import islice
from contextlib import contextmanager

infile = ...
outfile = ...

@contextmanager
def head(infile, numlines):
    with open(infile) as f:
        yield islice(f, numlines)

with open(outfile, "w") as f:
    with head(infile, 5) as lines:
        f.writelines(lines)