pyhexen/wad.py

   1
   2 import os
   3 import os.path
   4 import struct
   5
   6 class WadException(Exception):
   7         pass
   8
   9 class WadDesc:
  10         wadinfo_t = '<4sii'
  11         wadinfo_size = struct.calcsize (wadinfo_t)
  12         filelump_t = '<ii8s'
  13         filelump_size = struct.calcsize (filelump_t)
  14
  15 class WadReader:
  16         def __init__(self):
  17                 self.lumps = []
  18                 self.lumpcache = []
  19
  20         def __del__(self):
  21                 for (handle, pos, size, name) in self.lumps:
  22                         if not handle.closed:
  23                                 handle.close()
  24
  25         def AddFile(self, filename):
  26                 "Add a single file to the wad database"
  27                 (base, ext) = os.path.splitext (filename)
  28                 try:
  29                         fileHandle = open (filename, 'rb')
  30                 except IOError:
  31                         raise WadException ("Could not open '%s'" % filename)
  32
  33                 lumpinfo = []
  34                 if ext.lower() == '.wad':
  35                         # yay, a wad file.
  36                         (id, numLumps, infoTable) = struct.unpack (WadDesc.wadinfo_t, fileHandle.read(WadDesc.wadinfo_size))
  37
  38                         if id != 'IWAD' and id != 'PWAD':
  39                                 raise WadException("Wad file %s doesn't have IWAD or PWAD id" % filename)
  40
  41                         fileHandle.seek (infoTable)
  42
  43                         for i in range (0, numLumps):
  44                                 lumpinfo.append (struct.unpack (WadDesc.filelump_t, fileHandle.read(WadDesc.filelump_size)))
  45                 else:
  46                         # single file, boring.
  47                         size = os.stat (filename).st_size
  48                         lumpinfo.append ((0, size, base.upper()[:8]))
  49
  50                 # Dump the info into our table of accumulated crap.
  51                 for (pos, length, name) in lumpinfo:
  52                         self.lumps.append ((fileHandle, pos, length, name.rstrip('\x00')))
  53                         self.lumpcache.append ((False, ''))
  54
  55         def InitMultipleFiles (self, filenames):
  56                 "Add multiple files to the wad database"
  57                 for f in filenames:
  58                         self.AddFile (f)
  59
  60         def InitFile (self, filename):
  61                 "Add a single file to the wad database"
  62                 self.AddFile (f)
  63
  64         def NumLumps (self):
  65                 "Return the number of lumps in the wad database"
  66                 return len (self.lumps)
  67
  68         def CheckNumForName (self, name):
  69                 "Check if a specific lump name exists in the wad database, returning its index or -1 if it was not found"
  70                 searchName = name.upper()[:8]
  71                 for (x, y) in enumerate (reversed(self.lumps)):
  72                         (handle, pos, length, name) = y
  73                         if searchName == name:
  74                                 return len(self.lumps) - x - 1
  75                 return -1
  76
  77         def GetNumForName (self, name):
  78                 "Get the lump index of a specific lump name, or throw an exception if it was not found"
  79                 index = self.CheckNumForName (name)
  80                 if index == -1:
  81                         raise WadException ("%s not found!" % name)
  82                 return index
  83
  84         def LumpLength (self, lump):
  85                 "Return the size of a specific lump index"
  86                 if lump < 0 or lump > len(self.lumps):
  87                         raise WadException ('Invalid lump number %d requested' % lump)
  88                 (handle, pos, length, name) = self.lumps[lump]
  89                 return length
  90
  91         def ReadLump (self, lump):
  92                 "Return the data associated with a specific lump index"
  93                 if lump < 0 or lump > len(self.lumps):
  94                         raise WadException ('Invalid lump number %d requested' % lump)
  95                 (handle, pos, length, name) = self.lumps[lump]
  96                 handle.seek (pos)
  97                 return handle.read (length)
  98
  99         def CacheLumpNum (self, lump, tag=None):
 100                 "Cache and return the data associated with a specific lump index"
 101                 if lump < 0 or lump > len(self.lumps):
 102                         raise WadException ('Invalid lump number %d requested' % lump)
 103                 (cached, data) = self.lumpcache[lump]
 104                 if not cached:
 105                         data = self.ReadLump (lump)
 106                         self.lumpcache[lump] = (True, data)
 107                 return data
 108
 109         def CacheLumpName (self, name, tag=None):
 110                 "Cache and return the data associated with a specific lump name"
 111                 return self.CacheLumpNum (self.GetNumForName (name), tag)
 112