@@ -15,31 +15,29 @@ def __init__(self, filename):
15
15
16
16
def write_block (self , data , block_number ):
17
17
filename = self .BLOCK_FILENAME_FORMAT .format (block_number )
18
- file = open (filename , 'w' )
19
- file .write (data )
20
- file .close ()
18
+ with open (filename , 'w' ) as file :
19
+ file .write (data )
21
20
self .block_filenames .append (filename )
22
21
23
22
def get_block_filenames (self ):
24
23
return self .block_filenames
25
24
26
25
def split (self , block_size , sort_key = None ):
27
- file = open (self .filename , 'r' )
28
26
i = 0
27
+ with open (self .filename ) as file :
28
+ while True :
29
+ lines = file .readlines (block_size )
29
30
30
- while True :
31
- lines = file . readlines ( block_size )
31
+ if lines == [] :
32
+ break
32
33
33
- if lines == []:
34
- break
34
+ if sort_key is None :
35
+ lines .sort ()
36
+ else :
37
+ lines .sort (key = sort_key )
35
38
36
- if sort_key is None :
37
- lines .sort ()
38
- else :
39
- lines .sort (key = sort_key )
40
-
41
- self .write_block ('' .join (lines ), i )
42
- i += 1
39
+ self .write_block ('' .join (lines ), i )
40
+ i += 1
43
41
44
42
def cleanup (self ):
45
43
map (lambda f : os .remove (f ), self .block_filenames )
@@ -74,6 +72,7 @@ def refresh(self):
74
72
75
73
if self .buffers [i ] == '' :
76
74
self .empty .add (i )
75
+ self .files [i ].close ()
77
76
78
77
if len (self .empty ) == self .num_buffers :
79
78
return False
@@ -92,12 +91,11 @@ def __init__(self, merge_strategy):
92
91
self .merge_strategy = merge_strategy
93
92
94
93
def merge (self , filenames , outfilename , buffer_size ):
95
- outfile = open (outfilename , 'w' , buffer_size )
96
94
buffers = FilesArray (self .get_file_handles (filenames , buffer_size ))
97
-
98
- while buffers .refresh ():
99
- min_index = self .merge_strategy .select (buffers .get_dict ())
100
- outfile .write (buffers .unshift (min_index ))
95
+ with open ( outfilename , 'w' , buffer_size ) as outfile :
96
+ while buffers .refresh ():
97
+ min_index = self .merge_strategy .select (buffers .get_dict ())
98
+ outfile .write (buffers .unshift (min_index ))
101
99
102
100
def get_file_handles (self , filenames , buffer_size ):
103
101
files = {}
0 commit comments