comparison pytouhou/formats/pbg3.py @ 204:88361534c77e

Add some documentation (argh, so much left to document!)
author Thibaut Girka <thib@sitedethib.com>
date Tue, 01 Nov 2011 13:46:03 +0100
parents ac2e5e1c2c3c
children b5c7369abd7c
comparison
equal deleted inserted replaced
203:df8b2ab54639 204:88361534c77e
10 ## but WITHOUT ANY WARRANTY; without even the implied warranty of 10 ## but WITHOUT ANY WARRANTY; without even the implied warranty of
11 ## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 11 ## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12 ## GNU General Public License for more details. 12 ## GNU General Public License for more details.
13 ## 13 ##
14 14
15 """PBG3 archive files handling.
16
17 This module provides classes for handling the PBG3 file format.
18 The PBG3 format is the archive format used by Touhou: EoSD.
19
20 PBG3 files are merely a bitstream composed of a header,
21 a file table, and LZSS-compressed files.
22 """
23
24 from collections import namedtuple
25
15 from pytouhou.utils.bitstream import BitStream 26 from pytouhou.utils.bitstream import BitStream
16 import pytouhou.utils.lzss as lzss 27 import pytouhou.utils.lzss as lzss
17 28
18 from pytouhou.utils.helpers import get_logger 29 from pytouhou.utils.helpers import get_logger
19 30
20 logger = get_logger(__name__) 31 logger = get_logger(__name__)
21 32
22 33
23 class PBG3BitStream(BitStream): 34 class PBG3BitStream(BitStream):
35 """Helper class to handle strings and integers in PBG3 bitstreams."""
36
24 def read_int(self): 37 def read_int(self):
38 """Read an integer from the bitstream.
39
40 Integers have variable sizes. They begin with a two-bit value indicating
41 the number of (non-aligned) bytes to read.
42 """
43
25 size = self.read(2) 44 size = self.read(2)
26 return self.read((size + 1) * 8) 45 return self.read((size + 1) * 8)
27 46
28 47
29 def read_string(self, maxsize): 48 def read_string(self, maxsize):
49 """Read a string from the bitstream.
50
51 Strings are stored as standard NULL-termianted sequences of bytes.
52 The only catch is that they are not byte-aligned.
53 """
54
30 string = [] 55 string = []
31 for i in range(maxsize): 56 for i in range(maxsize):
32 byte = self.read(8) 57 byte = self.read(8)
33 if byte == 0: 58 if byte == 0:
34 break 59 break
35 string.append(byte) 60 string.append(byte)
36 return ''.join(chr(byte) for byte in string) 61 return ''.join(chr(byte) for byte in string)
37 62
38 63
39 64
65 PBG3Entry = namedtuple('PBG3Entry', 'unknown1 unknown2 checksum offset size')
66
67
68
40 class PBG3(object): 69 class PBG3(object):
41 def __init__(self, entries, bitstream=None): 70 """Handle PBG3 archive files.
42 self.entries = entries 71
72 PBG3 is a file archive format used in Touhou 6: EoSD.
73 This class provides a representation of such files, as well as functions to
74 read and extract files from a PBG3 archive.
75
76 Instance variables:
77 entries -- list of PBG3Entry objects describing files present in the archive
78 bitstream -- PBG3BitStream object
79 """
80
81 def __init__(self, entries=None, bitstream=None):
82 self.entries = entries or []
43 self.bitstream = bitstream #TODO 83 self.bitstream = bitstream #TODO
44 84
45 85
46 def __enter__(self): 86 def __enter__(self):
47 return self 87 return self
51 return self.bitstream.__exit__(type, value, traceback) 91 return self.bitstream.__exit__(type, value, traceback)
52 92
53 93
54 @classmethod 94 @classmethod
55 def read(cls, file): 95 def read(cls, file):
96 """Read a PBG3 file.
97
98 Raise an exception if the file is invalid.
99 Return a PBG3 instance otherwise.
100 """
101
56 magic = file.read(4) 102 magic = file.read(4)
57 if magic != b'PBG3': 103 if magic != b'PBG3':
58 raise Exception #TODO 104 raise Exception #TODO
59 105
60 bitstream = PBG3BitStream(file) 106 bitstream = PBG3BitStream(file)
68 unknown2 = bitstream.read_int() 114 unknown2 = bitstream.read_int()
69 checksum = bitstream.read_int() # Checksum of *compressed data* 115 checksum = bitstream.read_int() # Checksum of *compressed data*
70 offset = bitstream.read_int() 116 offset = bitstream.read_int()
71 size = bitstream.read_int() 117 size = bitstream.read_int()
72 name = bitstream.read_string(255).decode('ascii') 118 name = bitstream.read_string(255).decode('ascii')
73 entries[name] = (unknown1, unknown2, checksum, offset, size) 119 entries[name] = PBG3Entry(unknown1, unknown2, checksum, offset, size)
74 120
75 return PBG3(entries, bitstream) 121 return PBG3(entries, bitstream)
76 122
77 123
78 def list_files(self): 124 def list_files(self):
125 """List files present in the archive."""
79 return self.entries.keys() 126 return self.entries.keys()
80 127
81 128
82 def extract(self, filename, check=False): 129 def extract(self, filename, check=False):
130 """Extract a given file.
131
132 If “filename” is in the archive, extract it and return its contents.
133 Otherwise, raise an exception.
134
135 By default, the checksum of the file won't be verified,
136 you can however force the verification using the “check” argument.
137 """
138
83 unkwn1, unkwn2, checksum, offset, size = self.entries[filename] 139 unkwn1, unkwn2, checksum, offset, size = self.entries[filename]
84 self.bitstream.seek(offset) 140 self.bitstream.seek(offset)
85 data = lzss.decompress(self.bitstream, size) 141 data = lzss.decompress(self.bitstream, size)
86 if check: 142 if check:
87 # Checking the checksum 143 # Verify the checksum
88 compressed_size = self.bitstream.io.tell() - offset 144 compressed_size = self.bitstream.io.tell() - offset
89 self.bitstream.seek(offset) 145 self.bitstream.seek(offset)
90 value = 0 146 value = 0
91 for c in self.bitstream.io.read(compressed_size): 147 for c in self.bitstream.io.read(compressed_size):
92 value += ord(c) 148 value += ord(c)