1 #!/usr/bin/env python 2 3 """ 4 Generic file access. 5 6 Copyright (C) 2009, 2010, 2011 Paul Boddie <paul@boddie.org.uk> 7 8 This program is free software; you can redistribute it and/or modify it under 9 the terms of the GNU General Public License as published by the Free Software 10 Foundation; either version 3 of the License, or (at your option) any later 11 version. 12 13 This program is distributed in the hope that it will be useful, but WITHOUT ANY 14 WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A 15 PARTICULAR PURPOSE. See the GNU General Public License for more details. 16 17 You should have received a copy of the GNU General Public License along 18 with this program. If not, see <http://www.gnu.org/licenses/>. 19 """ 20 21 from iixr.data import * 22 from array import array 23 import zlib 24 25 # Constants. 26 27 class File: 28 29 "A basic file abstraction." 30 31 def __init__(self, f): 32 self.f = f 33 self.data = array('B') # master buffer 34 self.record = array('B') # record buffer 35 self.reset() 36 37 def reset(self): 38 39 "To be used to reset the state of the reader or writer between records." 40 41 pass 42 43 def seek(self, offset): 44 self.f.seek(offset) 45 self.reset() 46 47 def rewind(self): 48 self.f.seek(0) 49 self.reset() 50 51 def close(self): 52 if self.f is not None: 53 self.f.close() 54 self.f = None 55 56 class FileWriter(File): 57 58 "Writing basic data types to files." 59 60 def tell(self): 61 # NOTE: Will not be accurate within the current record. 62 return self.f.tell() + len(self.data) 63 64 def begin_record(self): 65 pass 66 67 def end_record(self): 68 if self.record: 69 vint_to_array(len(self.record), self.data) 70 self.data += self.record 71 self.record = array('B') 72 73 def write_number(self, number): 74 75 "Write 'number' to the file using a variable length encoding." 76 77 vint_to_array(number, self.record) 78 79 def write_numbers(self, numbers): 80 81 "Write 'numbers' to the file using a variable length encoding." 82 83 for number in numbers: 84 vint_to_array(number, self.record) 85 86 def write_string(self, s, compress=0): 87 88 """ 89 Write 's' to the file, recording its length and compressing the string 90 if 'compress' is set to a true value. 91 """ 92 93 # Convert Unicode objects to strings. 94 95 if isinstance(s, unicode): 96 s = s.encode("utf-8") 97 98 # Compress the string if requested. 99 100 if compress: 101 cs = zlib.compress(s) 102 103 # Take any shorter than the original. 104 105 if len(cs) < len(s): 106 flag = "z" 107 s = cs 108 else: 109 flag = "-" 110 111 else: 112 flag = "" 113 114 # Write the length of the data before the data itself. 115 116 length = len(s) 117 self.record.fromstring("".join([flag, vint(length), s])) 118 119 def write_sequence_value(self, value): 120 sequence_to_array(value, self.record) 121 122 def write_sequence_values(self, values): 123 vint_to_array(len(values), self.record) 124 for value in values: 125 self.write_sequence_value(value) 126 127 def write_delta_sequence(self, values): 128 convert_sequence(values, get_subtractor(values[0])) 129 self.write_sequence_values(values) 130 131 def write_monotonic_sequence(self, values): 132 convert_sequence(values, get_monotonic_subtractor(values[0])) 133 self.write_sequence_values(values) 134 135 def flush(self): 136 if self.f is not None: 137 self.end_record() 138 self.data.tofile(self.f) 139 self.data = array('B') 140 141 def close(self): 142 self.flush() 143 File.close(self) 144 145 class FileReader(File): 146 147 "Reading basic data types from files." 148 149 def begin_record(self): 150 self.record = array('B') 151 self.start = 0 152 try: 153 size = self.read_number_from_file() 154 self.record.fromfile(self.f, size) 155 except EOFError: 156 pass 157 158 def end_record(self): 159 pass 160 161 def read_number_from_file(self): 162 163 "Read a number from the file." 164 165 # Read each byte, adding it to the number. 166 167 f = self.f 168 a = array('B') 169 fromfile = a.fromfile 170 171 fromfile(f, 1) 172 csd = a[-1] 173 if csd < 128: 174 return csd 175 else: 176 while csd & 128: 177 fromfile(f, 1) 178 csd = a[-1] 179 return vint_from_array(a) 180 181 def read_number(self): 182 183 "Read a number from the current record." 184 185 n, self.start = vint_from_array_start(self.record, self.start) 186 return n 187 188 def read_string(self, decompress=0): 189 190 """ 191 Read a string from the current record, decompressing the stored data if 192 'decompress' is set to a true value. 193 """ 194 195 # Decompress the data if requested. 196 197 if decompress: 198 flag = chr(self.record[self.start]) 199 self.start += 1 200 else: 201 flag = "-" 202 203 length = self.read_number() 204 start = self.start 205 self.start += length 206 s = self.record[start:self.start].tostring() 207 208 # Perform decompression if applicable. 209 210 if flag == "z": 211 s = zlib.decompress(s) 212 213 # Convert strings to Unicode objects. 214 215 return unicode(s, "utf-8") 216 217 def read_sequence_value(self): 218 value, self.start = sequence_from_array(self.record, self.start) 219 return value 220 221 def read_sequences(self): 222 values = [] 223 length = self.read_number() 224 i = 0 225 while i < length: 226 values.append(self.read_sequence_value()) 227 i += 1 228 return values 229 230 def read_delta_sequence(self): 231 values = self.read_sequences() 232 convert_sequence(values, get_adder(values[0])) 233 return values 234 235 def read_monotonic_sequence(self): 236 values = self.read_sequences() 237 convert_sequence(values, get_monotonic_adder(values[0])) 238 return values 239 240 # vim: tabstop=4 expandtab shiftwidth=4