|
7 | 7 | """Create and extract SigMF archives.""" |
8 | 8 |
|
9 | 9 | import io |
10 | | -import os |
11 | 10 | import shutil |
12 | 11 | import tarfile |
13 | 12 | import tempfile |
| 13 | +from pathlib import Path |
14 | 14 |
|
15 | 15 | from .error import SigMFFileError |
16 | 16 |
|
|
21 | 21 |
|
22 | 22 |
|
23 | 23 | class SigMFArchive: |
24 | | - """Archive a SigMFFile. |
| 24 | + """ |
| 25 | + Archive a SigMFFile |
25 | 26 |
|
26 | 27 | A `.sigmf` file must include both valid metadata and data. |
27 | 28 | If `self.data_file` is not set or the requested output file |
28 | | - is not writable, raise `SigMFFileError`. |
29 | | -
|
30 | | - Parameters: |
31 | | -
|
32 | | - sigmffile -- A SigMFFile object with valid metadata and data_file |
33 | | -
|
34 | | - name -- path to archive file to create. If file exists, overwrite. |
35 | | - If `name` doesn't end in .sigmf, it will be appended. |
36 | | - For example: if `name` == "/tmp/archive1", then the |
37 | | - following archive will be created: |
38 | | - /tmp/archive1.sigmf |
39 | | - - archive1/ |
40 | | - - archive1.sigmf-meta |
41 | | - - archive1.sigmf-data |
42 | | -
|
43 | | - fileobj -- If `fileobj` is specified, it is used as an alternative to |
44 | | - a file object opened in binary mode for `name`. It is |
45 | | - supposed to be at position 0. `name` is not required, but |
46 | | - if specified will be used to determine the directory and |
47 | | - file names within the archive. `fileobj` won't be closed. |
48 | | - For example: if `name` == "archive1" and fileobj is given, |
49 | | - a tar archive will be written to fileobj with the |
50 | | - following structure: |
51 | | - - archive1/ |
52 | | - - archive1.sigmf-meta |
53 | | - - archive1.sigmf-data |
| 29 | + is not writable, raises `SigMFFileError`. |
| 30 | +
|
| 31 | + Parameters |
| 32 | + ---------- |
| 33 | +
|
| 34 | + sigmffile : SigMFFile |
| 35 | + A SigMFFile object with valid metadata and data_file. |
| 36 | +
|
| 37 | + name : PathLike | str | bytes |
| 38 | + Path to archive file to create. If file exists, overwrite. |
| 39 | + If `name` doesn't end in .sigmf, it will be appended. |
| 40 | + For example: if `name` == "/tmp/archive1", then the |
| 41 | + following archive will be created: |
| 42 | + /tmp/archive1.sigmf |
| 43 | + - archive1/ |
| 44 | + - archive1.sigmf-meta |
| 45 | + - archive1.sigmf-data |
| 46 | +
|
| 47 | + fileobj : BufferedWriter |
| 48 | + If `fileobj` is specified, it is used as an alternative to |
| 49 | + a file object opened in binary mode for `name`. It is |
| 50 | + supposed to be at position 0. `name` is not required, but |
| 51 | + if specified will be used to determine the directory and |
| 52 | + file names within the archive. `fileobj` won't be closed. |
| 53 | + For example: if `name` == "archive1" and fileobj is given, |
| 54 | + a tar archive will be written to fileobj with the |
| 55 | + following structure: |
| 56 | + - archive1/ |
| 57 | + - archive1.sigmf-meta |
| 58 | + - archive1.sigmf-data |
54 | 59 | """ |
55 | 60 |
|
56 | 61 | def __init__(self, sigmffile, name=None, fileobj=None): |
| 62 | + is_buffer = fileobj is not None |
57 | 63 | self.sigmffile = sigmffile |
58 | | - self.name = name |
59 | | - self.fileobj = fileobj |
60 | | - |
61 | | - self._check_input() |
| 64 | + self.path, arcname, fileobj = self._resolve(name, fileobj) |
62 | 65 |
|
63 | | - archive_name = self._get_archive_name() |
64 | | - sigmf_fileobj = self._get_output_fileobj() |
65 | | - sigmf_archive = tarfile.TarFile(mode="w", fileobj=sigmf_fileobj, format=tarfile.PAX_FORMAT) |
66 | | - tmpdir = tempfile.mkdtemp() |
67 | | - sigmf_md_filename = archive_name + SIGMF_METADATA_EXT |
68 | | - sigmf_md_path = os.path.join(tmpdir, sigmf_md_filename) |
69 | | - sigmf_data_filename = archive_name + SIGMF_DATASET_EXT |
70 | | - sigmf_data_path = os.path.join(tmpdir, sigmf_data_filename) |
| 66 | + self._ensure_data_file_set() |
| 67 | + self._validate() |
71 | 68 |
|
72 | | - with open(sigmf_md_path, "w") as mdfile: |
73 | | - self.sigmffile.dump(mdfile, pretty=True) |
| 69 | + tar = tarfile.TarFile(mode="w", fileobj=fileobj, format=tarfile.PAX_FORMAT) |
| 70 | + tmpdir = Path(tempfile.mkdtemp()) |
| 71 | + meta_path = tmpdir / (arcname + SIGMF_METADATA_EXT) |
| 72 | + data_path = tmpdir / (arcname + SIGMF_DATASET_EXT) |
74 | 73 |
|
| 74 | + # write files |
| 75 | + with open(meta_path, "w") as handle: |
| 76 | + self.sigmffile.dump(handle) |
75 | 77 | if isinstance(self.sigmffile.data_buffer, io.BytesIO): |
76 | | - self.sigmffile.data_file = sigmf_data_path |
77 | | - with open(sigmf_data_path, "wb") as f: |
78 | | - f.write(self.sigmffile.data_buffer.getbuffer()) |
| 78 | + # write data buffer to archive |
| 79 | + self.sigmffile.data_file = data_path |
| 80 | + with open(data_path, "wb") as handle: |
| 81 | + handle.write(self.sigmffile.data_buffer.getbuffer()) |
79 | 82 | else: |
80 | | - shutil.copy(self.sigmffile.data_file, sigmf_data_path) |
81 | | - |
82 | | - def chmod(tarinfo): |
83 | | - if tarinfo.isdir(): |
84 | | - tarinfo.mode = 0o755 # dwrxw-rw-r |
85 | | - else: |
86 | | - tarinfo.mode = 0o644 # -wr-r--r-- |
87 | | - return tarinfo |
88 | | - |
89 | | - sigmf_archive.add(tmpdir, arcname=archive_name, filter=chmod) |
90 | | - sigmf_archive.close() |
91 | | - if not fileobj: |
92 | | - sigmf_fileobj.close() |
93 | | - |
| 83 | + # copy data to archive |
| 84 | + shutil.copy(self.sigmffile.data_file, data_path) |
| 85 | + tar.add(tmpdir, arcname=arcname, filter=self.chmod) |
| 86 | + # close files & remove tmpdir |
| 87 | + tar.close() |
| 88 | + if not is_buffer: |
| 89 | + # only close fileobj if we aren't working w/a buffer |
| 90 | + fileobj.close() |
94 | 91 | shutil.rmtree(tmpdir) |
95 | 92 |
|
96 | | - self.path = sigmf_archive.name |
97 | | - |
98 | | - def _check_input(self): |
99 | | - self._ensure_name_has_correct_extension() |
100 | | - self._ensure_data_file_set() |
101 | | - self._validate_sigmffile_metadata() |
102 | | - |
103 | | - def _ensure_name_has_correct_extension(self): |
104 | | - name = self.name |
105 | | - if name is None: |
106 | | - return |
107 | | - |
108 | | - has_extension = "." in name |
109 | | - has_correct_extension = name.endswith(SIGMF_ARCHIVE_EXT) |
110 | | - if has_extension and not has_correct_extension: |
111 | | - apparent_ext = os.path.splitext(name)[-1] |
112 | | - err = "extension {} != {}".format(apparent_ext, SIGMF_ARCHIVE_EXT) |
113 | | - raise SigMFFileError(err) |
114 | | - |
115 | | - self.name = name if has_correct_extension else name + SIGMF_ARCHIVE_EXT |
| 93 | + @staticmethod |
| 94 | + def chmod(tarinfo: tarfile.TarInfo): |
| 95 | + """permission filter for writing tar files""" |
| 96 | + if tarinfo.isdir(): |
| 97 | + tarinfo.mode = 0o755 # dwrxw-rw-r |
| 98 | + else: |
| 99 | + tarinfo.mode = 0o644 # -wr-r--r-- |
| 100 | + return tarinfo |
116 | 101 |
|
117 | 102 | def _ensure_data_file_set(self): |
118 | 103 | if not self.sigmffile.data_file and not isinstance(self.sigmffile.data_buffer, io.BytesIO): |
119 | | - err = "no data file - use `set_data_file`" |
120 | | - raise SigMFFileError(err) |
| 104 | + raise SigMFFileError("No data file in SigMFFile; use `set_data_file` before archiving.") |
121 | 105 |
|
122 | | - def _validate_sigmffile_metadata(self): |
| 106 | + def _validate(self): |
123 | 107 | self.sigmffile.validate() |
124 | 108 |
|
125 | | - def _get_archive_name(self): |
126 | | - if self.fileobj and not self.name: |
127 | | - pathname = self.fileobj.name |
128 | | - else: |
129 | | - pathname = self.name |
130 | | - |
131 | | - filename = os.path.split(pathname)[-1] |
132 | | - archive_name, archive_ext = os.path.splitext(filename) |
133 | | - return archive_name |
134 | | - |
135 | | - def _get_output_fileobj(self): |
136 | | - try: |
137 | | - fileobj = self._get_open_fileobj() |
138 | | - except: |
139 | | - if self.fileobj: |
140 | | - err = "fileobj {!r} is not byte-writable".format(self.fileobj) |
141 | | - else: |
142 | | - err = "can't open {!r} for writing".format(self.name) |
143 | | - |
144 | | - raise SigMFFileError(err) |
145 | | - |
146 | | - return fileobj |
147 | | - |
148 | | - def _get_open_fileobj(self): |
149 | | - if self.fileobj: |
150 | | - fileobj = self.fileobj |
151 | | - fileobj.write(bytes()) # force exception if not byte-writable |
| 109 | + def _resolve(self, name, fileobj): |
| 110 | + """ |
| 111 | + Resolve both (name, fileobj) into (path, arcname, fileobj) given either or both. |
| 112 | +
|
| 113 | + Returns |
| 114 | + ------- |
| 115 | + path : PathLike |
| 116 | + Path of the archive file. |
| 117 | + arcname : str |
| 118 | + Name of the sigmf object within the archive. |
| 119 | + fileobj : BufferedWriter |
| 120 | + Open file handle object. |
| 121 | + """ |
| 122 | + if fileobj: |
| 123 | + try: |
| 124 | + # exception if not byte-writable |
| 125 | + fileobj.write(bytes()) |
| 126 | + # exception if no name property of handle |
| 127 | + path = Path(fileobj.name) |
| 128 | + if not name: |
| 129 | + arcname = path.stem |
| 130 | + else: |
| 131 | + arcname = name |
| 132 | + except io.UnsupportedOperation: |
| 133 | + raise SigMFFileError(f"fileobj {fileobj} is not byte-writable.") |
| 134 | + except AttributeError: |
| 135 | + raise SigMFFileError(f"fileobj {fileobj} is invalid.") |
| 136 | + elif name: |
| 137 | + path = Path(name) |
| 138 | + # ensure name has correct suffix if it exists |
| 139 | + if path.suffix == "": |
| 140 | + # add extension if none was given |
| 141 | + path = path.with_suffix(SIGMF_ARCHIVE_EXT) |
| 142 | + elif path.suffix != SIGMF_ARCHIVE_EXT: |
| 143 | + # ensure suffix is correct |
| 144 | + raise SigMFFileError(f"Invalid extension ({path.suffix} != {SIGMF_ARCHIVE_EXT}).") |
| 145 | + arcname = path.stem |
| 146 | + |
| 147 | + try: |
| 148 | + fileobj = open(path, "wb") |
| 149 | + except (OSError, IOError): |
| 150 | + raise SigMFFileError(f"Can't open {name} for writing.") |
152 | 151 | else: |
153 | | - fileobj = open(self.name, "wb") |
| 152 | + raise SigMFFileError("Either `name` or `fileobj` needs to be defined.") |
154 | 153 |
|
155 | | - return fileobj |
| 154 | + return path, arcname, fileobj |
0 commit comments