wadls - list all map files within a WAD/PK3/ZIP

wadls.py

#!/usr/bin/python

import os, sys, zipfile, tempfile

# wadls (pronounced "waddles", thx joshthenesnerd) - list all maps in a wad/zip/pk3

# requires omgifol module, set path to it here or in env variable
OMG_PATH = os.environ.get('OMG_PATH', None) or '/home/jpl/projects/wadsmoosh'
sys.path.append(OMG_PATH)
import omg

maps_found = 0

def arg_valid(arg):
    return os.path.exists(sys.argv[1]) and \
        os.path.splitext(sys.argv[1])[1].lower() in ['.wad', '.zip', '.pk3']

def list_maps_in_wad(wad_filename, display_name=None):
    global maps_found
    wad = omg.WAD()
    try:
        wad.from_file(wad_filename)
    except Exception as e:
        print("Couldn't read from %s: %s" % (wad_filename, e))
        return
    map_names = []
    maps = wad.maps.find('*')
    if len(maps) == 0:
        return
    # handle wads containing UDMF map data specially
    if len(maps) == 1 and maps[0] == 'ZNODES':
        print('  %s is a UDMF map' % display_name)
        maps_found += 1
        return
    print('%s:' % (display_name if display_name else wad_filename))
    for map_name in maps:
        map_names.append('  ' + map_name)
        maps_found += 1
    map_names.sort()
    print('\n'.join(map_names))

def log_maps_found():
    print('%s %s found in %s' % (maps_found,
                                 ['maps', 'map'][maps_found == 1],
                                 filename))

if __name__ == '__main__':
    if len(sys.argv) < 2 or not arg_valid(sys.argv[1]):
        print('Please specify a WAD, ZIP, or PK3 file.')
        sys.exit()
    filename = sys.argv[1]
    # if it's a WAD, just list any maps
    if filename.lower().endswith('.wad'):
        list_maps_in_wad(filename)
        log_maps_found()
        sys.exit()
    # if it's a ZIP/PK3, find any WADs inside it
    z = zipfile.ZipFile(filename)
    for f in z.namelist():
        if not f.lower().endswith('.wad'):
            continue
        # write zipped file to temp location
        temp_file = tempfile.NamedTemporaryFile()
        zipped_file = z.open(f)
        temp_file.write(zipped_file.read())
        # flush so omg can read it
        temp_file.flush()
        list_maps_in_wad(temp_file.name, f)
        temp_file.close()
    z.close()
    log_maps_found()