|
1 | 1 | """fileio.py: OMF Writer and Reader for serializing to and from .omf files""" |
2 | 2 |
|
| 3 | +from __future__ import annotations |
| 4 | + |
3 | 5 | import json |
4 | 6 | import struct |
5 | 7 | import uuid |
6 | 8 |
|
7 | | -from six import string_types |
8 | | - |
9 | 9 | from omf.base import UidModel |
10 | 10 | from omf.fileio.geoh5 import GeoH5Writer |
11 | 11 |
|
@@ -36,13 +36,14 @@ class OMFWriter: |
36 | 36 | in the binary blob. |
37 | 37 | """ |
38 | 38 |
|
39 | | - def __init__(self, project, fname): |
| 39 | + def __init__(self, project: UidModel, fname: str, compression: int = 5): |
40 | 40 | """Project serialization is performed on OMFWriter init |
41 | 41 |
|
42 | 42 | Binary data is written during project serialization |
43 | 43 | """ |
| 44 | + |
44 | 45 | if fname.endswith("geoh5"): |
45 | | - GeoH5Writer(project, fname) |
| 46 | + GeoH5Writer(project, fname, compression=compression) |
46 | 47 | else: |
47 | 48 | if not fname.endswith(".omf"): |
48 | 49 | fname = fname + ".omf" |
@@ -100,7 +101,7 @@ class OMFReader: |
100 | 101 |
|
101 | 102 | def __init__(self, fopen): |
102 | 103 | if isinstance(fopen, str): |
103 | | - fopen = open(fopen, "rb") |
| 104 | + fopen = open(fopen, "rb") # pylint: disable=R1732 |
104 | 105 | self._fopen = fopen |
105 | 106 | fopen.seek(0, 0) |
106 | 107 | self._uid, self._json_start = self.read_header() |
@@ -157,8 +158,8 @@ def read_header(self): |
157 | 158 | file_version = file_version[0 : len(__version__)] |
158 | 159 | if file_version != __version__: |
159 | 160 | raise ValueError( |
160 | | - "Version mismatch: file version {fv!r}, " |
161 | | - "reader version {rv!r}".format(fv=file_version, rv=__version__) |
| 161 | + f"Version mismatch: file version {file_version!r}, " |
| 162 | + f"reader version {__version__!r}" |
162 | 163 | ) |
163 | 164 | uid = uuid.UUID(bytes=struct.unpack("<16s", self._fopen.read(16))[0]) |
164 | 165 | json_start = struct.unpack("<Q", self._fopen.read(8))[0] |
|
0 commit comments