diff --git a/geomagio/Util.py b/geomagio/Util.py index dff464a09bf60fe9c1c38d03a448f1984687d3ed..d1cb0a73b1ffc4c164961cddedb8e2b70f973802 100644 --- a/geomagio/Util.py +++ b/geomagio/Util.py @@ -1,6 +1,7 @@ import numpy import os from obspy.core import Stats, Trace +from obspy import UTCDateTime from io import BytesIO import json import fcntl @@ -203,7 +204,28 @@ def create_empty_trace(trace, channel): return Trace(numpy_data, stats) -def write_state_file(filename, data, directory=None): +def encode_utcdatetime(obj): + """ + Custom JSON encoder for dealing with UTCDateTime objects + """ + if isinstance(obj, UTCDateTime): + return str(obj) + raise TypeError( + f"Object of type '{obj.__class__.__name__}' is not JSON serializable" + ) + + +def decode_utcdatetime(dct): + """ + Custom JSON decoder for converting time fields back to UTCDateTime objects + """ + for key in ["start_time", "end_time", "starttime", "endtime"]: + if key in dct: + dct[key] = UTCDateTime(dct[key]) if dct[key] else None + return dct + + +def write_state_file(filename, data, directory=None, encoder=None): """ Writes data to a state file in a thread-safe manner. @@ -215,6 +237,8 @@ def write_state_file(filename, data, directory=None): The data to write to the file. This should be a Python object that can be serialized with json. directory: String The directory to write the file to. If not provided, the file will be written to the .cache directory in the current user's home directory. + encoder: function + Function to be given to json.dump's 'default' parameter. If not provided it will use a simple encoder that handles UTCDateTime objects. Returns: -------- @@ -228,6 +252,9 @@ def write_state_file(filename, data, directory=None): if directory is None: directory = os.path.join(os.path.expanduser("~"), ".cache", "geomag-algorithms") + if encoder is None: + encoder = encode_utcdatetime + # Create the directory if it doesn't exist try: os.makedirs(directory, exist_ok=True) @@ -241,7 +268,7 @@ def write_state_file(filename, data, directory=None): with open(filepath, "w") as f: try: fcntl.flock(f, fcntl.LOCK_EX) - json.dump(data, f) + json.dump(data, f, default=encoder) fcntl.flock(f, fcntl.LOCK_UN) except IOError as e: print(f"Error locking or writing to file: {e}") @@ -254,7 +281,7 @@ def write_state_file(filename, data, directory=None): raise -def read_state_file(filename, directory=None): +def read_state_file(filename, directory=None, decoder=None): """ Reads data from a state file in a thread-safe manner. @@ -263,6 +290,8 @@ def read_state_file(filename, directory=None): The name of the file to read from. directory: String The directory to read the file from. If not provided, the file will be read from the .cache directory in the current user's home directory. + encoder: function + Object hook function to be given to json.load. If not provided it will use a simple decoder that handles common start/end time fields. Returns: -------- @@ -277,13 +306,16 @@ def read_state_file(filename, directory=None): if directory is None: directory = os.path.join(os.path.expanduser("~"), ".cache", "geomag-algorithms") + if decoder is None: + decoder = decode_utcdatetime + filepath = os.path.join(directory, filename) try: with open(filepath, "r") as f: try: fcntl.flock(f, fcntl.LOCK_SH) - data = json.load(f) + data = json.load(f, object_hook=decoder) fcntl.flock(f, fcntl.LOCK_UN) return data except IOError as e: diff --git a/geomagio/metadata/instrument/InstrumentCalibrations.py b/geomagio/metadata/instrument/InstrumentCalibrations.py index a78f2db1bbf935952368a3925d45c6066231c788..1bdc33d792afe18c9e1f2afb2ed9789738cfff02 100644 --- a/geomagio/metadata/instrument/InstrumentCalibrations.py +++ b/geomagio/metadata/instrument/InstrumentCalibrations.py @@ -354,7 +354,10 @@ def get_instrument_calibrations( calibrations = read_state_file(filename=state_filename) else: - write_state_file(state_filename, calibrations) + try: + write_state_file(state_filename, calibrations) + except TypeError: + print("Write state file failed, object couldn't be encoded.") return [ c