gw: Clean up image-file handling. Create a helper Image base class.

scripts/greaseweazle/image/adf.py

@@ -7,26 +7,23 @@
 
 from greaseweazle import error
 import greaseweazle.codec.amiga.amigados as amigados
+from .image import Image
 
-class ADF:
+class ADF(Image):
 
     default_format = 'amiga.amigados'
 
     def __init__(self, start_cyl, nr_sides):
         error.check(nr_sides == 2, "ADF: Must be double-sided")
-        self.bitrate = 253
         self.sec_per_track = 11
         self.track_list = [None] * start_cyl
 
 
     @classmethod
-    def to_file(cls, start_cyl, nr_sides):
-        adf = cls(start_cyl, nr_sides)
-        return adf
-
+    def from_file(cls, name):
 
-    @classmethod
-    def from_file(cls, dat):
+        with open(name, "rb") as f:
+            dat = f.read()
 
         adf = cls(0, 2)
 
@@ -36,7 +33,6 @@ class ADF:
         if ncyl > 90:
             ncyl //= 2
             nsec *= 2
-            adf.bitrate *= 2
             adf.sec_per_track = nsec
 
         for i in range(ncyl*2):

scripts/greaseweazle/image/hfe.py

@@ -10,8 +10,9 @@ import struct
 from greaseweazle import error
 from greaseweazle.track import MasterTrack, RawTrack
 from bitarray import bitarray
+from .image import Image
 
-class HFE:
+class HFE(Image):
 
     def __init__(self, start_cyl, nr_sides):
         self.start_cyl = start_cyl
@@ -23,13 +24,10 @@ class HFE:
 
 
     @classmethod
-    def to_file(cls, start_cyl, nr_sides):
-        hfe = cls(start_cyl, nr_sides)
-        return hfe
-
+    def from_file(cls, name):
 
-    @classmethod
-    def from_file(cls, dat):
+        with open(name, "rb") as f:
+            dat = f.read()
 
         (sig, f_rev, nr_cyls, nr_sides, t_enc, bitrate,
          _, _, _, tlut_base) = struct.unpack("<8s4B2H2BH", dat[:20])

scripts/greaseweazle/image/image.py

@@ -0,0 +1,47 @@
+# greaseweazle/image/image.py
+#
+# Written & released by Keir Fraser <keir.xen@gmail.com>
+#
+# This is free and unencumbered software released into the public domain.
+# See the file COPYING for more details, or visit <http://unlicense.org>.
+
+import os
+
+from greaseweazle import error
+import greaseweazle.codec.amiga.amigados as amigados
+
+class Image:
+
+    read_only = False
+
+    ## Context manager for image objects created using .to_file()
+
+    def __enter__(self):
+        self.file = open(self.filename, "wb")
+        return self
+
+    def __exit__(self, type, value, tb):
+        try:
+            if type is None:
+                # No error: Normal writeout.
+                self.file.write(self.get_image())
+        finally:
+            # Always close the file.
+            self.file.close()
+        if type is not None:
+            # An error occurred: We remove the target file.
+            os.remove(self.filename)
+
+    ## Default .to_file() constructor
+    @classmethod
+    def to_file(cls, name, start_cyl, nr_sides):
+        error.check(not cls.read_only,
+                    "%s: Cannot create %s image files" % (name, cls.__name__))
+        obj = cls(start_cyl, nr_sides)
+        obj.filename = name
+        return obj
+
+
+# Local variables:
+# python-indent: 4
+# End:

scripts/greaseweazle/image/ipf.py

@@ -11,6 +11,7 @@ import ctypes as ct
 from bitarray import bitarray
 from greaseweazle.track import MasterTrack
 from greaseweazle import error
+from .image import Image
 
 class CapsDateTimeExt(ct.Structure):
     _pack_ = 1
@@ -98,7 +99,9 @@ class DI_LOCK:
     def_flags = (DENVAR | UPDATEFD | TYPE | OVLBIT | TRKBIT)
 
 
-class IPF:
+class IPF(Image):
+
+    read_only = True
 
     def __init__(self, start_cyl, nr_sides):
         self.lib = get_libcaps()
@@ -135,7 +138,7 @@ class IPF:
         return s
 
     @classmethod
-    def from_filename(cls, name):
+    def from_file(cls, name):
 
         ipf = cls(0, 0)
 

scripts/greaseweazle/image/scp.py

@@ -9,6 +9,7 @@ import struct, functools
 
 from greaseweazle import error
 from greaseweazle.flux import Flux
+from .image import Image
 
 class SCPOpts:
     """legacy_ss: Set to True to generate (incorrect) legacy single-sided
@@ -17,7 +18,7 @@ class SCPOpts:
     def __init__(self):
         self.legacy_ss = False
 
-class SCP:
+class SCP(Image):
 
     # 40MHz
     sample_freq = 40000000
@@ -30,13 +31,10 @@ class SCP:
 
 
     @classmethod
-    def to_file(cls, start_cyl, nr_sides):
-        scp = cls(start_cyl, nr_sides)
-        return scp
-
+    def from_file(cls, name):
 
-    @classmethod
-    def from_file(cls, dat):
+        with open(name, "rb") as f:
+            dat = f.read()
 
         header = struct.unpack("<3s9BI", dat[0:16])
         (sig, _, _, nr_revs, _, _, flags, _, single_sided, _, _) = header

scripts/greaseweazle/tools/read.py

@@ -19,10 +19,7 @@ from greaseweazle.flux import Flux
 
 
 def open_image(args, image_class):
-    error.check(hasattr(image_class, 'to_file'),
-                "%s: Cannot create %s image files"
-                % (args.file, image_class.__name__))
-    image = image_class.to_file(args.scyl, args.nr_sides)
+    image = image_class.to_file(args.file, args.scyl, args.nr_sides)
     if args.rate is not None:
         image.bitrate = args.rate
     for opt, val in args.file_opts.items():
@@ -120,9 +117,6 @@ def read_to_image(usb, args, image, decoder=None):
     if decoder is not None:
         print_summary(args, summary)
 
-    # Write the image file.
-    with open(args.file, "wb") as f:
-        f.write(image.get_image())
 
 def range_str(s, e):
     str = "%d" % s
@@ -179,9 +173,9 @@ def main(argv):
               (range_str(args.scyl, args.ecyl),
                range_str(0, args.nr_sides-1),
                args.revs))
-        image = open_image(args, image_class)
-        util.with_drive_selected(read_to_image, usb, args, image,
-                                 decoder=decoder)
+        with open_image(args, image_class) as image:
+            util.with_drive_selected(read_to_image, usb, args, image,
+                                     decoder=decoder)
     except USB.CmdError as err:
         print("Command Failed: %s" % err)
 

scripts/greaseweazle/tools/write.py

@@ -15,16 +15,10 @@ from greaseweazle.tools import util
 from greaseweazle import error
 from greaseweazle import usb as USB
 
-
 # Read and parse the image file.
 def open_image(args):
-    image_class = util.get_image_class(args.file)
-    if hasattr(image_class, 'from_filename'):
-        image = image_class.from_filename(args.file)
-    else:
-        with open(args.file, "rb") as f:
-            image = image_class.from_file(f.read())
-    return image
+    cls = util.get_image_class(args.file)
+    return cls.from_file(args.file)
 
 class Formatter:
     def __init__(self):