Add optional encoder/decoder params to json util functions

geomagio/Util.py

@@ -225,7 +225,7 @@ def decode_utcdatetime(dct):
     return dct
 
 
-def write_state_file(filename, data, directory=None):
+def write_state_file(filename, data, directory=None, encoder=None):
     """
     Writes data to a state file in a thread-safe manner.
 
@@ -237,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:
     --------
@@ -250,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)
@@ -263,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, default=encode_utcdatetime)
+                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}")
@@ -276,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.
 
@@ -285,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:
     --------
@@ -299,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, object_hook=decode_utcdatetime)
+                data = json.load(f, object_hook=decoder)
                 fcntl.flock(f, fcntl.LOCK_UN)
                 return data
             except IOError as e: