os.path converted pathlib.Path

sigmf/apps/convert_wav.py

@@ -9,72 +9,78 @@
 import argparse
 import getpass
 import logging
-import os
-import pathlib
 import tempfile
+from datetime import datetime, timezone
+from os import PathLike
+from pathlib import Path
 from typing import Optional
 
 from scipy.io import wavfile
 
-from .. import SigMFFile, __specification__
+from .. import SigMFFile
 from .. import __version__ as toolversion
-from .. import archive
-from ..utils import get_data_type_str, get_sigmf_iso8601_datetime_now
+from ..sigmffile import get_sigmf_filenames
+from ..utils import SIGMF_DATETIME_ISO8601_FMT, get_data_type_str
 
 log = logging.getLogger()
 
 
 def convert_wav(
-    input_wav_filename: str,
-    archive_filename: Optional[str],
-    start_datetime: Optional[str] = None,
+    wav_path: str,
+    out_path: Optional[str] = None,
     author: Optional[str] = None,
-):
+) -> PathLike:
     """
-    read a .wav and write a .sigmf archive
+    Read a wav and write a sigmf archive.
     """
-    input_path = pathlib.Path(input_wav_filename)
-    input_stem = input_path.stem
-    samp_rate, wav_data = wavfile.read(input_wav_filename)
+    wav_path = Path(wav_path)
+    wav_stem = wav_path.stem
+    samp_rate, wav_data = wavfile.read(wav_path)
 
     global_info = {
         SigMFFile.AUTHOR_KEY: getpass.getuser() if author is None else author,
         SigMFFile.DATATYPE_KEY: get_data_type_str(wav_data),
-        SigMFFile.DESCRIPTION_KEY: f"Converted from {input_wav_filename}",
+        SigMFFile.DESCRIPTION_KEY: f"converted from {wav_path.name}",
         SigMFFile.NUM_CHANNELS_KEY: 1 if len(wav_data.shape) < 2 else wav_data.shape[1],
-        SigMFFile.RECORDER_KEY: os.path.basename(__file__),
+        SigMFFile.RECORDER_KEY: "Official SigMF wav converter",
         SigMFFile.SAMPLE_RATE_KEY: samp_rate,
-        SigMFFile.VERSION_KEY: __specification__,
     }
 
-    if start_datetime is None:
-        start_datetime = get_sigmf_iso8601_datetime_now()
+    modify_time = wav_path.lstat().st_mtime
+    wav_datetime = datetime.fromtimestamp(modify_time, tz=timezone.utc)
 
     capture_info = {
         SigMFFile.START_INDEX_KEY: 0,
-        SigMFFile.DATETIME_KEY: start_datetime,
+        SigMFFile.DATETIME_KEY: wav_datetime.strftime(SIGMF_DATETIME_ISO8601_FMT),
     }
 
-    tmpdir = tempfile.mkdtemp()
-    sigmf_data_filename = input_stem + archive.SIGMF_DATASET_EXT
-    sigmf_data_path = os.path.join(tmpdir, sigmf_data_filename)
-    wav_data.tofile(sigmf_data_path)
+    temp_dir = Path(tempfile.mkdtemp())
+    if out_path is None:
+        # extension will be changed
+        out_path = Path(wav_stem)
+    else:
+        out_path = Path(out_path)
+    filenames = get_sigmf_filenames(out_path)
 
-    meta = SigMFFile(data_file=sigmf_data_path, global_info=global_info)
+    data_path = temp_dir / filenames["data_fn"]
+    wav_data.tofile(data_path)
+
+    meta = SigMFFile(data_file=data_path, global_info=global_info)
     meta.add_capture(0, metadata=capture_info)
+    log.debug("created %r", meta)
 
-    if archive_filename is None:
-        archive_filename = input_stem + archive.SIGMF_ARCHIVE_EXT
-    meta.tofile(archive_filename, toarchive=True)
-    return os.path.abspath(archive_filename)
+    arc_path = filenames["archive_fn"]
+    meta.tofile(arc_path, toarchive=True)
+    log.info("wrote %s", arc_path)
+    return arc_path
 
 
-def main():
+def main() -> None:
     """
     entry-point for sigmf_convert_wav
     """
-    parser = argparse.ArgumentParser(description="Convert .wav to .sigmf container.")
-    parser.add_argument("input", type=str, help="Wavfile path")
+    parser = argparse.ArgumentParser(description="Convert wav to sigmf archive.")
+    parser.add_argument("input", type=str, help="wav path")
     parser.add_argument("--author", type=str, default=None, help=f"set {SigMFFile.AUTHOR_KEY} metadata")
     parser.add_argument("-v", "--verbose", action="count", default=0)
     parser.add_argument("--version", action="version", version=f"%(prog)s v{toolversion}")
@@ -87,11 +93,10 @@ def main():
     }
     logging.basicConfig(level=level_lut[min(args.verbose, 2)])
 
-    out_fname = convert_wav(
-        input_wav_filename=args.input,
+    _ = convert_wav(
+        wav_path=args.input,
         author=args.author,
     )
-    log.info(f"Write {out_fname}")
 
 
 if __name__ == "__main__":

sigmf/archive.py

@@ -7,10 +7,10 @@
 """Create and extract SigMF archives."""
 
 import io
-import os
 import shutil
 import tarfile
 import tempfile
+from pathlib import Path
 
 from .error import SigMFFileError
 
@@ -21,135 +21,134 @@
 
 
 class SigMFArchive:
-    """Archive a SigMFFile.
+    """
+    Archive a SigMFFile
 
     A `.sigmf` file must include both valid metadata and data.
     If `self.data_file` is not set or the requested output file
-    is not writable, raise `SigMFFileError`.
-
-    Parameters:
-
-      sigmffile -- A SigMFFile object with valid metadata and data_file
-
-      name      -- path to archive file to create. If file exists, overwrite.
-                   If `name` doesn't end in .sigmf, it will be appended.
-                   For example: if `name` == "/tmp/archive1", then the
-                   following archive will be created:
-                       /tmp/archive1.sigmf
-                       - archive1/
-                         - archive1.sigmf-meta
-                         - archive1.sigmf-data
-
-      fileobj   -- If `fileobj` is specified, it is used as an alternative to
-                   a file object opened in binary mode for `name`. It is
-                   supposed to be at position 0. `name` is not required, but
-                   if specified will be used to determine the directory and
-                   file names within the archive. `fileobj` won't be closed.
-                   For example: if `name` == "archive1" and fileobj is given,
-                   a tar archive will be written to fileobj with the
-                   following structure:
-                       - archive1/
-                         - archive1.sigmf-meta
-                         - archive1.sigmf-data
+    is not writable, raises `SigMFFileError`.
+
+    Parameters
+    ----------
+
+    sigmffile : SigMFFile
+        A SigMFFile object with valid metadata and data_file.
+
+    name : PathLike | str | bytes
+        Path to archive file to create. If file exists, overwrite.
+        If `name` doesn't end in .sigmf, it will be appended.
+        For example: if `name` == "/tmp/archive1", then the
+        following archive will be created:
+            /tmp/archive1.sigmf
+            - archive1/
+                - archive1.sigmf-meta
+                - archive1.sigmf-data
+
+    fileobj : BufferedWriter
+        If `fileobj` is specified, it is used as an alternative to
+        a file object opened in binary mode for `name`. It is
+        supposed to be at position 0. `name` is not required, but
+        if specified will be used to determine the directory and
+        file names within the archive. `fileobj` won't be closed.
+        For example: if `name` == "archive1" and fileobj is given,
+        a tar archive will be written to fileobj with the
+        following structure:
+            - archive1/
+                - archive1.sigmf-meta
+                - archive1.sigmf-data
     """
 
     def __init__(self, sigmffile, name=None, fileobj=None):
+        is_buffer = fileobj is not None
         self.sigmffile = sigmffile
-        self.name = name
-        self.fileobj = fileobj
-
-        self._check_input()
+        self.path, arcname, fileobj = self._resolve(name, fileobj)
 
-        archive_name = self._get_archive_name()
-        sigmf_fileobj = self._get_output_fileobj()
-        sigmf_archive = tarfile.TarFile(mode="w", fileobj=sigmf_fileobj, format=tarfile.PAX_FORMAT)
-        tmpdir = tempfile.mkdtemp()
-        sigmf_md_filename = archive_name + SIGMF_METADATA_EXT
-        sigmf_md_path = os.path.join(tmpdir, sigmf_md_filename)
-        sigmf_data_filename = archive_name + SIGMF_DATASET_EXT
-        sigmf_data_path = os.path.join(tmpdir, sigmf_data_filename)
+        self._ensure_data_file_set()
+        self._validate()
 
-        with open(sigmf_md_path, "w") as mdfile:
-            self.sigmffile.dump(mdfile, pretty=True)
+        tar = tarfile.TarFile(mode="w", fileobj=fileobj, format=tarfile.PAX_FORMAT)
+        tmpdir = Path(tempfile.mkdtemp())
+        meta_path = tmpdir / (arcname + SIGMF_METADATA_EXT)
+        data_path = tmpdir / (arcname + SIGMF_DATASET_EXT)
 
+        # write files
+        with open(meta_path, "w") as handle:
+            self.sigmffile.dump(handle)
         if isinstance(self.sigmffile.data_buffer, io.BytesIO):
-            self.sigmffile.data_file = sigmf_data_path
-            with open(sigmf_data_path, "wb") as f:
-                f.write(self.sigmffile.data_buffer.getbuffer())
+            # write data buffer to archive
+            self.sigmffile.data_file = data_path
+            with open(data_path, "wb") as handle:
+                handle.write(self.sigmffile.data_buffer.getbuffer())
         else:
-            shutil.copy(self.sigmffile.data_file, sigmf_data_path)
-
-        def chmod(tarinfo):
-            if tarinfo.isdir():
-                tarinfo.mode = 0o755  # dwrxw-rw-r
-            else:
-                tarinfo.mode = 0o644  # -wr-r--r--
-            return tarinfo
-
-        sigmf_archive.add(tmpdir, arcname=archive_name, filter=chmod)
-        sigmf_archive.close()
-        if not fileobj:
-            sigmf_fileobj.close()
-
+            # copy data to archive
+            shutil.copy(self.sigmffile.data_file, data_path)
+        tar.add(tmpdir, arcname=arcname, filter=self.chmod)
+        # close files & remove tmpdir
+        tar.close()
+        if not is_buffer:
+            # only close fileobj if we aren't working w/a buffer
+            fileobj.close()
         shutil.rmtree(tmpdir)
 
-        self.path = sigmf_archive.name
-
-    def _check_input(self):
-        self._ensure_name_has_correct_extension()
-        self._ensure_data_file_set()
-        self._validate_sigmffile_metadata()
-
-    def _ensure_name_has_correct_extension(self):
-        name = self.name
-        if name is None:
-            return
-
-        has_extension = "." in name
-        has_correct_extension = name.endswith(SIGMF_ARCHIVE_EXT)
-        if has_extension and not has_correct_extension:
-            apparent_ext = os.path.splitext(name)[-1]
-            err = "extension {} != {}".format(apparent_ext, SIGMF_ARCHIVE_EXT)
-            raise SigMFFileError(err)
-
-        self.name = name if has_correct_extension else name + SIGMF_ARCHIVE_EXT
+    @staticmethod
+    def chmod(tarinfo: tarfile.TarInfo):
+        """permission filter for writing tar files"""
+        if tarinfo.isdir():
+            tarinfo.mode = 0o755  # dwrxw-rw-r
+        else:
+            tarinfo.mode = 0o644  # -wr-r--r--
+        return tarinfo
 
     def _ensure_data_file_set(self):
         if not self.sigmffile.data_file and not isinstance(self.sigmffile.data_buffer, io.BytesIO):
-            err = "no data file - use `set_data_file`"
-            raise SigMFFileError(err)
+            raise SigMFFileError("No data file in SigMFFile; use `set_data_file` before archiving.")
 
-    def _validate_sigmffile_metadata(self):
+    def _validate(self):
         self.sigmffile.validate()
 
-    def _get_archive_name(self):
-        if self.fileobj and not self.name:
-            pathname = self.fileobj.name
-        else:
-            pathname = self.name
-
-        filename = os.path.split(pathname)[-1]
-        archive_name, archive_ext = os.path.splitext(filename)
-        return archive_name
-
-    def _get_output_fileobj(self):
-        try:
-            fileobj = self._get_open_fileobj()
-        except:
-            if self.fileobj:
-                err = "fileobj {!r} is not byte-writable".format(self.fileobj)
-            else:
-                err = "can't open {!r} for writing".format(self.name)
-
-            raise SigMFFileError(err)
-
-        return fileobj
-
-    def _get_open_fileobj(self):
-        if self.fileobj:
-            fileobj = self.fileobj
-            fileobj.write(bytes())  # force exception if not byte-writable
+    def _resolve(self, name, fileobj):
+        """
+        Resolve both (name, fileobj) into (path, arcname, fileobj) given either or both.
+
+        Returns
+        -------
+        path : PathLike
+            Path of the archive file.
+        arcname : str
+            Name of the sigmf object within the archive.
+        fileobj : BufferedWriter
+            Open file handle object.
+        """
+        if fileobj:
+            try:
+                # exception if not byte-writable
+                fileobj.write(bytes())
+                # exception if no name property of handle
+                path = Path(fileobj.name)
+                if not name:
+                    arcname = path.stem
+                else:
+                    arcname = name
+            except io.UnsupportedOperation:
+                raise SigMFFileError(f"fileobj {fileobj} is not byte-writable.")
+            except AttributeError:
+                raise SigMFFileError(f"fileobj {fileobj} is invalid.")
+        elif name:
+            path = Path(name)
+            # ensure name has correct suffix if it exists
+            if path.suffix == "":
+                # add extension if none was given
+                path = path.with_suffix(SIGMF_ARCHIVE_EXT)
+            elif path.suffix != SIGMF_ARCHIVE_EXT:
+                # ensure suffix is correct
+                raise SigMFFileError(f"Invalid extension ({path.suffix} != {SIGMF_ARCHIVE_EXT}).")
+            arcname = path.stem
+
+            try:
+                fileobj = open(path, "wb")
+            except (OSError, IOError):
+                raise SigMFFileError(f"Can't open {name} for writing.")
         else:
-            fileobj = open(self.name, "wb")
+            raise SigMFFileError("Either `name` or `fileobj` needs to be defined.")
 
-        return fileobj
+        return path, arcname, fileobj