Reading first lines of bz2 files in python

Here is a fully working example that includes writing and reading a test file that is much smaller than your 10000 lines. Its nice to have working examples in questions so we can test easily.

import bz2
import itertools
import codecs

file = "file.bz2"
file_10000 = "file.txt"

# write test file with 9 lines
with bz2.BZ2File(file, "w") as fp:
    fp.write('\n'.join('123456789'))

# the original script using BZ2File ... and 3 lines for test
# ...and fixing bugs:
#     1) it only writes 9999 instead of 10000
#     2) files don't do writerow
#     3) close the files

output_file = codecs.open(file_10000,'w+','utf-8')

source_file = bz2.BZ2File(file, "r")
count = 0
for line in source_file:
    count += 1
    if count <= 3:
       output_file.write(line)
source_file.close()
output_file.close()

# show what you got
print('---- Test 1 ----')
print(repr(open(file_10000).read()))   

A more efficient way to do it is to break out of the for loop after reading the lines you want. you can even leverage iterators to thin up the code like so:

# a faster way to read first 3 lines
with bz2.BZ2File(file) as source_file,\
        codecs.open(file_10000,'w+','utf-8') as output_file:
    output_file.writelines(itertools.islice(source_file, 3))

# show what you got
print('---- Test 2 ----')
print(repr(open(file_10000).read()))   

This is definitely a simpler way of doing it than the other answer, but it would be an easy way to do so in both Python2/3. Also, it would short-circuit if you don't have >= 10,000 lines.

from bz2 import BZ2File as bzopen

# writing to a file
with bzopen("file.bz2", "w") as bzfout:
    for i in range(123456):
        bzfout.write(b"%i\n" % i)

# reading a bz2 archive
with bzopen("file.bz2", "r") as bzfin:
    """ Handle lines here """
    lines = []
    for i, line in enumerate(bzfin):
        if i == 10000: break
        lines.append(line.rstrip())

print(lines)

Tags:

Python

Bz2