A simple attempt at working out what files are in and out of a mod, and if its broken

dep.py

import os, sys, string
import xml.dom.minidom as minidom
from struct import unpack
from itertools import chain

class File:
    """a file (type and path)"""
    MAP = "map"
    SCENARIO = "scenario"
    FRACTION = "faction"
    TILESET = "tileset"
    TEXTURE = "texture"
    TECH_TREE = "tech-tree"
    UNIT = "unit"
    FACTION = "faction"
    MODEL = "model"
    SOUND = "sound"
    PARTICLE = "particle"
    UPGRADE = "upgrade"
    RESOURCE = "resource"
    LANGUAGE = "language"
    def __init__(self,mod,typ,path):
        self.mod = mod
        self.typ = typ
        self.path = path
        self.references = set()
        self.referenced_by = set()
        self.broken = False
        self._filesize = 0
        self.virtual = path.startswith("!")
        if not self.virtual and not path.startswith(mod.base_folder):
            self.error("external dependency not yet supported")
    def error(self,*args):
        self.broken = True
        broken = self.mod.broken
        if self not in broken:
            broken[self] = []
        for prev in broken[self]:
            if prev == args:
                break
        else:
            broken[self].append(args)
    def subpath(self,r):
        if r.startswith("/"):
            r = r[1:]
        return os.path.join(os.path.split(self.path)[0],r)
    def filesize(self):
        if not self.broken and self._filesize == 0:
            self._filesize = os.path.getsize(self.path)
        return self._filesize
    def __repr__(self):
        folder = self.mod.base_folder if self.virtual else os.path.split(self.path)[0]
        references = ", ".join(["%s %s"%(r.typ,r.path[1:] if r.virtual else os.path.relpath(r.path,folder)) for r in self.references])
        referenced_by = ", ".join(["%s %s"%(r.typ,r.path[1:] if r.virtual else os.path.relpath(r.path,folder)) for r in self.referenced_by])
        return "(%s%s %s%s%s%s)"%(self.typ,
            " BROKEN" if self.broken else "",
            self.path if self.virtual else os.path.relpath(self.path,self.mod.base_folder),
            " (references: %s)"%references if references is not "" else "",
            " (referenced by: %s)"%referenced_by if referenced_by is not "" else "",
            " %s"%fmt_bytes(self.filesize()) if not self.broken else "")
    def sortorder(self,other):
        # don't override __cmp__ because we want to be hashable
        assert isinstance(other,File)
        if other.typ == self.typ:
            return -cmp(other.path,self.path)
        return -cmp(other.typ,self.typ)

class Files:
    def __init__(self,mod):
        self.mod = mod
        self.files = {}
        self.ignored = set()
        self.typ = {}
    def ref(self,typ,path,referenced_by):
        assert isinstance(referenced_by,File)
        f = self.add(typ,path)
        f.referenced_by.add(referenced_by)
        referenced_by.references.add(f)
        return f
    def add(self,typ,path):
        # tidy up path
        unsafe_chars = None
        if not path.startswith("!"):
            path = os.path.abspath(path)
            for i,ch in enumerate(os.path.relpath(path,self.mod.base_folder)):
                if not ((ch in string.lowercase) or\
                    (ch in string.digits) or\
                    (ch in "/\\._")):
                    ### unsafe_chars = "(%d,%d,%c)"%(i+1,ord(ch),ch)
                    break
        if path in self.files:
            f = self.files[path]
            assert f.typ == typ
        else:
            f = self.files[path] = File(self.mod,typ,path)
        if unsafe_chars is not None:
            f.error("Path contains unsafe characters",unsafe_chars)
        if typ not in self.typ:
            self.typ[typ] = set()
        self.typ[typ].add(f)
        return f
        
class FilterExt:
    def __init__(self,*ext):
        assert len(ext) > 0
        self.ext = ext
    def __call__(self,f):
        ext = os.path.splitext(f)[1].lower()
        return ext in self.ext

class Mod:
    def __init__(self,base_folder):
        self.base_folder = os.path.abspath(base_folder)
        self.maps = set()
        self.scenarios = set()
        self.factions = set()
        self.tilesets = set()
        self.tech_trees = set()
        self.files = Files(self)
        self.broken = {}
        self._init_maps()
        self._init_tilesets()
        self._init_tech_trees()
        self._init_scenarios()
        self._check_exists()
        self._init_ignored()
    def _init_maps(self):
        for f in self._listdir("maps",os.path.isfile,FilterExt(".mgm",".gbm")):
            self.files.add(File.MAP,f)
            self.maps.add(os.path.splitext(os.path.split(f)[1])[0])
    def _init_tilesets(self):
        for f in self._listdir("tilesets",os.path.isdir):
            name = os.path.split(f)[1]
            self.tilesets.add(name)
            f = os.path.join(f,"%s.xml"%name)
            f = self.files.add(File.TILESET,f)
            xml = self._init_xml(f)
            if xml is None:
                return
            assert xml.documentElement.tagName == "tileset", f
    def _init_tech_trees(self):
        for f in self._listdir("techs",os.path.isdir):
            name = os.path.split(f)[1]
            self.tech_trees.add(name)
            self.files.add(File.TECH_TREE,os.path.join(f,"%s.xml"%name))
            for f in self._listdir("techs/%s/factions"%name,os.path.isdir):
                self._init_faction(f)
            for f in self._listdir("techs/%s/resources"%name,os.path.isdir):
                self._init_resource(f)
    def _init_faction(self,f):
        name = os.path.split(f)[1]
        self.factions.add(name)
        self.files.add(File.FACTION,os.path.join(f,"%s.xml"%name))
        for unit in self._listdir(os.path.join(f,"units"),os.path.isdir):
            self._init_unit(unit)
        for upgrade in self._listdir(os.path.join(f,"upgrades"),os.path.isdir):
            self._init_upgrade(upgrade)
    def _init_xml(self,f):
        if not os.path.isfile(f.path):
            return
        xml = minidom.parse(f.path)
        def extract(x,attr,typ):
            path = None
            try:
                if x.attributes["value"].value == "true":
                    path = x.attribute[attr].value
            except:
                try:
                    path = x.attributes[attr].value
                except:
                    pass
            if path is not None:
                return self.files.ref(typ,f.subpath(path),f)
        for sound in chain(xml.getElementsByTagName("sound"),
            xml.getElementsByTagName("sound-file"),
            xml.getElementsByTagName("music")):
            extract(sound,"path",File.SOUND)
        for image in chain(xml.getElementsByTagName("image"),
            xml.getElementsByTagName("texture"),
            xml.getElementsByTagName("image-cancel")):
            extract(image,"path",File.TEXTURE)
        for image in xml.getElementsByTagName("meeting-point"):
            extract(image,"image-path",File.TEXTURE)
        for model in chain(xml.getElementsByTagName("animation"),
            xml.getElementsByTagName("model")):
            model = extract(model,"path",File.MODEL)
            if model is not None:
                self._init_model(model)
        for particle in xml.getElementsByTagName("particle"):
            particle = extract(particle,"path",File.PARTICLE)
            if particle is not None:
                self._init_particle(particle)
        for sound in xml.getElementsByTagName("ambient-sounds"):
            for sound in [c for c in sound.childNodes if c.nodeType == c.ELEMENT_NODE]:
                extract(sound,"path",File.SOUND)
        return xml
    def _init_unit(self,f):
        name = os.path.split(f)[1]
        f = self.files.add(File.UNIT,os.path.join(f,"%s.xml"%name))
        f.name = name
        xml = self._init_xml(f)
        if xml is None:
            return f
        assert xml.documentElement.tagName == "unit", f
        for unit in chain(xml.documentElement.getElementsByTagName("unit"),
            xml.documentElement.getElementsByTagName("produced-unit")):
            unit = unit.attributes["name"].value
            self.files.ref(File.UNIT,os.path.join(f.path,"../../%s/%s.xml"%(unit,unit)),f)
        return f
    def _init_upgrade(self,f):
        name = os.path.split(f)[1]
        f = self.files.add(File.UPGRADE,os.path.join(f,"%s.xml"%name))
        f.name = name
        xml = self._init_xml(f)
        if xml is None:
            return
        assert xml.documentElement.tagName == "upgrade", f
    def _init_resource(self,f):
        name = os.path.split(f)[1]
        f = self.files.add(File.RESOURCE,os.path.join(f,"%s.xml"%name))
        f.name = name
        xml = self._init_xml(f)
        if xml is None:
            return
        assert xml.documentElement.tagName == "resource", f
    def _init_particle(self,f):
        if hasattr(f,"inited") and f.inited:
            return
        f.inited = True
        xml = self._init_xml(f)
        if xml is None:
            return
        assert xml.documentElement.tagName in ["projectile-particle-system",
            "splash-particle-system","unit-particle-system","particle-system"],\
            xml.documentElement.tagName
    def _init_scenarios(self):
        for f in self._listdir("scenarios",os.path.isdir):
            name = os.path.split(f)[1]
            self.scenarios.add(name)
            f = os.path.join(f,"%s.xml"%name)
            f = self.files.add(File.SCENARIO,f)
            xml = minidom.parse(f.path)
            assert xml.documentElement.tagName == "scenario", f
            factions = set()
            for player in xml.getElementsByTagName("player"):
                try:
                    if player.attributes["control"].value != "closed":
                        factions.add(player.attributes["faction"].value)
                except Exception as e:
                    print player.toxml(),e
            for faction in factions.difference(self.factions):
                f.error("References missing faction",faction)
            for m in xml.getElementsByTagName("map"):
                m = m.attributes["value"].value
                if m not in self.maps:
                    f.error("References missing map",m)
            for tileset in xml.getElementsByTagName("tileset"):
                tileset = tileset.attributes["value"].value
                if tileset not in self.tilesets:
                    f.error("References missing tileset",tileset)
            for tech_tree in xml.getElementsByTagName("tech-tree"):
                tech_tree = tech_tree.attributes["value"].value
                if tech_tree not in self.tech_trees:
                    f.error("References missing tech tree",tech_tree)
            for lng in self._listdir("scenarios/%s"%name,os.path.isfile,FilterExt(".lng")):
                if os.path.split(lng)[1].startswith(name):
                    self.files.add(File.LANGUAGE,lng)
    def _init_model(self,model):
        if hasattr(model,"inited") and model.inited:
            return
        model.inited = True
        try:
            f = open(model.path,"rb")
            if not f.read(3) == "G3D":
                model.error("not a valid G3D model")
                return
            ver = ord(f.read(1))
            def uint16():
                return unpack("<H",f.read(2))[0]
            def uint32():
                return unpack("<L",f.read(4))[0]
            if ver == 3:
                meshCount = uint32()
                for mesh in xrange(meshCount):
                    vertexFrameCount = uint32()
                    normalFrameCount = uint32()
                    texCoordFrameCount = uint32()
                    colorFrameCount = uint32()
                    pointCount = uint32()
                    indexCount = uint32()
                    properties = uint32()
                    texture = f.read(64)
                    has_texture = 0 == (properties & 1)
                    if has_texture:
                        texture = texture[:texture.find('\0')]
                        self.files.ref(File.TEXTURE,model.subpath(texture),model)
                    f.read(12*vertexFrameCount*pointCount)
                    f.read(12*vertexFrameCount*pointCount)
                    if has_texture:
                        f.read(8*texCoordFrameCount*pointCount)
                    f.read(16)
                    f.read(16*(colorFrameCount-1))
                    f.read(4*indexCount)
            elif ver == 4:
                meshCount = uint16()
                if ord(f.read(1)) != 0:
                    model.error("not mtMorphMesh!")
                    return
                for mesh in xrange(meshCount):
                    f.read(64) # meshName
                    frameCount = uint32()
                    vertexCount = uint32()
                    indexCount = uint32()
                    f.read(8*4)
                    properties = uint32()
                    textures = uint32()
                    for t in xrange(5):
                        if ((1 << t) & textures) != 0:
                            texture = f.read(64)
                            texture = texture[:texture.find('\0')]
                            self.files.ref(File.TEXTURE,model.subpath(texture),model)
                    f.read(12*frameCount*vertexCount*2)
                    if textures != 0:
                        f.read(8*vertexCount)
                    f.read(4*indexCount)
            else:
                model.error("Unsupported G3D version"+ver)
        except Exception as e:
            model.error("Error reading G3D file",e)
    def _listdir(self,path,*filters):
        path = os.path.join(self.base_folder,path)
        try:
            return [f for f in map(lambda x: os.path.join(path,x),os.listdir(path)) if all([ftr(f) for ftr in filters])]
        except OSError as e:
            if e.errno == 2: # file not found
                return []
            raise
    def _check_exists(self):
        for f in self.files.files.values():
            if not os.path.isfile(f.path):
                path = os.path.split(f.path)[1].lower()
                candidate = None
                try:
                    for c in os.listdir(os.path.dirname(f.path)):
                        if c.lower() == path:
                            assert candidate is None,"WTF?"
                            candidate = path
                except:
                    pass
                if candidate is not None:
                    f.path = candidate
                    f.error("Must rename file to lowercase")
                else:
                    f.error("File does not exist")
    def _init_ignored(self):
        for folder in os.walk(self.base_folder):
            for f in folder[2]:
                f = os.path.join(folder[0],f)
                if f not in self.files.files:
                    self.files.ignored.add(f)
                    
def fmt_bytes(b):
    for m in ["B","KB","MB","GB"]:
        if b < 1024:
            return "%1.1f %s"%(b,m)
        b /= 1024.
        
def sum_type(mod,typ):
    if typ in mod.files.typ:
       label = string.capitalize(typ)
       typ = mod.files.typ[typ]
       print "=== %ss:"%label,len(typ),fmt_bytes(sum(f.filesize() for f in typ)),"==="
            
if __name__ == "__main__":
    if len(sys.argv) != 2:
        print "Usage: python",sys.argv[0],"[mod_root_dir]"
        sys.exit(1)
    mod = Mod(sys.argv[1])
    included = 0
    include_count = 0
    for f in sorted(mod.files.files.values(),lambda x,y: x.sortorder(y)):
        if f.broken:
            continue
        include_count += 1
        included += f.filesize()
        print "Including",os.path.relpath(f.path,mod.base_folder),fmt_bytes(f.filesize())
    if len(mod.files.ignored) > 0:
        ignored = 0
        print "=== The following",len(mod.files.ignored),"files are ignored ==="
        for f in sorted(mod.files.ignored):
            try:
                filesize = os.path.getsize(f)
                ignored += filesize
                print "Ignoring",os.path.relpath(f,mod.base_folder), fmt_bytes(os.path.getsize(f))
            except:
                pass
        print "=== Ignored:",len(mod.files.ignored),fmt_bytes(ignored),"==="
    print "=== Included:",include_count,fmt_bytes(included),"==="
    sum_type(mod,File.MODEL)
    sum_type(mod,File.TEXTURE)
    sum_type(mod,File.PARTICLE)
    sum_type(mod,File.SOUND)
    sum_type(mod,File.UNIT)
    sum_type(mod,File.UPGRADE)
    sum_type(mod,File.RESOURCE)
    sum_type(mod,File.LANGUAGE)
    if len(mod.factions) > 0:
        print "=== Fractions:",", ".join(mod.factions),"==="
    if len(mod.maps) > 0:
        print "=== Maps:",", ".join(mod.maps),"==="
    if len(mod.tilesets) > 0:
        print "=== Tile Sets:",", ".join(mod.tilesets),"==="
    if len(mod.tech_trees) > 0:
        print "=== Tech Trees:",", ".join(mod.tech_trees),"==="
    if len(mod.scenarios) > 0:
        print "=== Scenarios:",", ".join(mod.scenarios),"==="
    if len(mod.broken) > 0:
        print "=== Mod check failed ==="
        for f,reason in sorted(mod.broken.items(),lambda x,y: x[0].sortorder(y[0])):
            try:
                print f,", ".join(str(r) for r in reason)
            except Exception as e:
                print "ERROR",f
                print "REASON",reason
                raise
    print "=== Done ==="