scrview.py: ZX Spectrum SCREEN$ viewer & trdtool.py: TRD disk image tool

scrview.py

#!/usr/bin/env python

"""scrview.py: ZX Spectrum SCREEN$ viewer"""

import sys
import Image
from array import array

class ZXScreen:
    WIDTH = 256
    HEIGHT = 192

    def __init__(self):
        self.bitmap = array('B')
        self.attributes = array('B')

    def load(self, filename):
        with open(filename, 'rb') as f:
            self.bitmap.fromfile(f, 6144)
            self.attributes.fromfile(f, 768)

    def get_pixel_address(self, x, y):
        y76 = y & 0b11000000 # third of screen
        y53 = y & 0b00111000
        y20 = y & 0b00000111
        address = (y76 << 5) + (y20 << 8) + (y53 << 2) + (x >> 3)
        return address

    def get_attribute_address(self, x, y):
        y73 = y & 0b11111000
        address = (y73 << 2) + (x >> 3)
        return address

    def get_byte(self, x, y):
        return self.bitmap[ self.get_pixel_address(x,y) ]

    def get_attribute(self, x, y):
        return self.attributes[ self.get_attribute_address(x,y) ]

    def convert(self):
        img = Image.new('RGB', (ZXScreen.WIDTH, ZXScreen.HEIGHT), 'white')
        pixels = img.load()
        for y in xrange(ZXScreen.HEIGHT):
            for col in xrange(ZXScreen.WIDTH >> 3):
                x = col << 3
                byte = self.get_byte(x, y)
                attr = self.get_attribute(x, y)
                ink = attr & 0b0111
                paper = (attr >> 3) & 0b0111
                bright = (attr >> 6) & 1
                val = 0xcd if not bright else 0xff
                for bit in xrange(8):
                    bit_is_set = (byte >> (7 - bit)) & 1
                    color = ink if bit_is_set else paper
                    rgb = tuple(val * (color >> i & 1) for i in (1,2,0))
                    pixels[x + bit, y] = rgb
        return img

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print "Usage: %s filename.scr" % sys.argv[0]
        sys.exit(2)

    scale = 4
    screen = ZXScreen()
    screen.load(sys.argv[1])
    img = screen.convert()
    img = img.resize((ZXScreen.WIDTH*scale, ZXScreen.HEIGHT*scale), Image.NEAREST)
    img.show()

trdtool.py

#!/usr/bin/env python

"""ZX Spectrum TRD tool"""

import sys
import os
from struct import *
from array import array
from collections import namedtuple

class TRDImage:
    FILE_RECORD_FORMAT = "=8scHHBBB"
    DISK_RECORD_FORMAT = "=223sHBBBBHBH9sBB8s3s"

    MAX_FILES = 128
    SECTOR_SIZE = 256

    FileInfo = namedtuple(
        'FileInfo',
        'name ext start length sector_count sector track'
    )

    DiskInfo = namedtuple(
        'DiskInfo',
        'buffer dcu_sectors next_free_sector next_free_track disk_type \
        file_count max_sectors sectors_per_track zero1 blank9 zero2 \
        deleted_file_count disk_label zero3'
    )

    reclen = calcsize(FILE_RECORD_FORMAT)
    dreclen = calcsize(DISK_RECORD_FORMAT)

    def open(self, filename):
        self.filename = filename
        self.read_directory()

    def get_file_info_by_index(self, idx):
        fi = self.FileInfo(*unpack_from(self.FILE_RECORD_FORMAT, self.directory, idx * self.reclen))
        return fi

    def iterate_files(self):
        for idx in xrange(self.MAX_FILES):
            fi = self.get_file_info_by_index(idx)
            if fi.name[0] == "\x00": break
            yield fi

    def format_filename(self, ff):
        return "{}.{}".format(ff.name.strip(), ff.ext.strip())

    def get_file_info_by_name(self, name):
        for fi in self.iterate_files():
            filename = self.format_filename(fi)
            if name == filename: return fi
        return None

    def read_directory(self):
        self.directory = array('B')
        self.diskinfo = array('B')

        with open(self.filename, 'rb') as f:
            self.directory.fromfile(f, self.reclen * self.MAX_FILES) # track 0, sectors 0..7
            self.diskinfo.fromfile(f, self.dreclen) # track 0, sector 8

        self.di = self.DiskInfo(*unpack(self.DISK_RECORD_FORMAT, self.diskinfo))

    def view(self, greppable=False, hex=False):
        if not greppable:
            disk_types = {
                0x16: '80 track, 2 side',
                0x17: '40 track, 2 side',
                0x18: '80 track, 1 side',
                0x19: '40 track, 1 side'
            }

            di = self.di
            dt = disk_types[di.disk_type] if di.disk_type in disk_types else 'Unknown'
            dl = di.disk_label.strip()

            print
            print "Listing TRD Image: %s" % self.filename
            print "Media type: %s, %d sectors per track" % (dt, di.sectors_per_track)
            print "Capacity: %d sectors (%d bytes)" % (di.max_sectors, di.max_sectors * self.SECTOR_SIZE)
            print "Label: %s" % (dl if dl else 'none')
            print
            print "Filename     Start  Length  Sectors  Track  Sector"
            print "----------   -----  ------  -------  -----  ------"
            if not hex:
                fmt = "{fullname:11}  {start:5d}  {length:6d}  {sector_count:7d}  {track:5d}  {sector:6d}"
            else:
                fmt = "{fullname:11}  #{start:04X}   #{length:04X}  {sector_count:7d}  {track:5d}  {sector:6d}"
        else:
            fmt = "-rw{x}rw-rw- 1 nobody nobody {length:-8d} 01-01-1980 00:00 {fullname}"

        count = 0
        for fi in self.iterate_files():
            count += 1
            name = self.format_filename(fi)
            name = "%s.%s" % (fi.name.strip(), fi.ext.strip())
            print fmt.format(
                name=fi.name.strip(), ext=fi.ext.strip(), fullname=name,
                start=fi.start, length=fi.length, sector_count=fi.sector_count,
                track=fi.track, sector=fi.sector, x=('x' if fi.ext=='B' else '-')
            )

        if not greppable:
            print "----------   -----  ------  -------  -----  ------"
            print count, "file(s),", di.deleted_file_count, "deleted file(s)\n"


    def get_file_by_name(self, name, dst):
        fi = self.get_file_info_by_name(name)
        ofs = (fi.sector + fi.track * self.di.sectors_per_track) * self.SECTOR_SIZE
        buffer = array('B')
        with open(self.filename, 'rb') as f:
            f.seek(ofs)
            buffer.fromfile(f, fi.length) # todo
        with open(dst, 'wb') as f:
            buffer.tofile(f)


if __name__ == '__main__':
    argc = len(sys.argv)

    if argc == 1:
        print "Usage: %s [command] disk.trd [boot.B] [/tmp/boot.B]" % sys.argv[0]
        print
        sys.exit(2)

    elif argc == 2:
        cmd = 'view'
        filename = sys.argv[1]

    else:
        cmd = sys.argv[1]
        filename = sys.argv[2]

    trd = TRDImage()
    trd.open(filename)

    if cmd == 'list':
        trd.view(greppable=True)

    elif cmd == 'view':
        trd.view(greppable=False)

    elif cmd == 'copyout':
        src = sys.argv[3]
        dst = sys.argv[4]
        trd.get_file_by_name(src, dst)

    elif cmd == 'copyin':
        pass

    elif cmd == 'run':
        pass

    else:
        print "ERROR: Unknown command: %s" % cmd
        sys.exit(2)