From 3025e63a504e3d62f591cb49bdff3aff98e2581b Mon Sep 17 00:00:00 2001
From: mivanit <mivanits@umich.edu>
Date: Mon, 13 May 2024 22:24:49 -0600
Subject: [PATCH] added `write_only_format` option and docstring to
 `JsonSerializer`

---
 muutils/json_serialize/json_serialize.py      | 34 +++++++++++++++++--
 .../test_sdc_defaults.py                      | 15 ++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/muutils/json_serialize/json_serialize.py b/muutils/json_serialize/json_serialize.py
index 4e9e2158..a857765f 100644
--- a/muutils/json_serialize/json_serialize.py
+++ b/muutils/json_serialize/json_serialize.py
@@ -206,7 +206,30 @@ def _serialize_override_serialize_func(
 
 
 class JsonSerializer:
-    """Json serialization class (holds configs)"""
+    """Json serialization class (holds configs)
+
+    # Parameters:
+        - `array_mode : ArrayMode`
+        how to write arrays
+        (defaults to `"array_list_meta"`)
+        - `error_mode : ErrorMode`
+        what to do when we can't serialize an object (will use repr as fallback if "ignore" or "warn")
+        (defaults to `"except"`)
+        - `handlers_pre : MonoTuple[SerializerHandler]`
+        handlers to use before the default handlers
+        (defaults to `tuple()`)
+        - `handlers_default : MonoTuple[SerializerHandler]`
+        default handlers to use
+        (defaults to `DEFAULT_HANDLERS`)
+        - `write_only_format : bool`
+        changes "__format__" keys in output to "__write_format__" (when you want to serialize something in a way that zanj won't try to recover the object when loading)
+        (defaults to `False`)
+
+    # Raises:
+        - `ValueError`: on init, if `args` is not empty
+        - `SerializationException`: on `json_serialize()`, if any error occurs when trying to serialize an object and `error_mode` is set to `"except"`
+
+    """
 
     def __init__(
         self,
@@ -215,6 +238,7 @@ def __init__(
         error_mode: ErrorMode = "except",
         handlers_pre: MonoTuple[SerializerHandler] = tuple(),
         handlers_default: MonoTuple[SerializerHandler] = DEFAULT_HANDLERS,
+        write_only_format: bool = False,
     ):
         if len(args) > 0:
             raise ValueError(
@@ -223,6 +247,7 @@ def __init__(
 
         self.array_mode: ArrayMode = array_mode
         self.error_mode: ErrorMode = error_mode
+        self.write_only_format: bool = write_only_format
         # join up the handlers
         self.handlers: MonoTuple[SerializerHandler] = tuple(handlers_pre) + tuple(
             handlers_default
@@ -236,7 +261,12 @@ def json_serialize(
         try:
             for handler in self.handlers:
                 if handler.check(self, obj, path):
-                    return handler.serialize_func(self, obj, path)
+                    output: JSONitem = handler.serialize_func(self, obj, path)
+                    if self.write_only_format:
+                        if isinstance(output, dict) and "__format__" in output:
+                            new_fmt: JSONitem = output.pop("__format__")
+                            output["__write_format__"] = new_fmt
+                    return output
 
             raise ValueError(f"no handler found for object with {type(obj) = }")
 
diff --git a/tests/unit/json_serialize/serializable_dataclass/test_sdc_defaults.py b/tests/unit/json_serialize/serializable_dataclass/test_sdc_defaults.py
index 01443e08..1272fc46 100644
--- a/tests/unit/json_serialize/serializable_dataclass/test_sdc_defaults.py
+++ b/tests/unit/json_serialize/serializable_dataclass/test_sdc_defaults.py
@@ -1,4 +1,5 @@
 from muutils.json_serialize import (
+    JsonSerializer,
     SerializableDataclass,
     serializable_dataclass,
     serializable_field,
@@ -27,6 +28,20 @@ def test_sdc_empty():
     assert recovered == instance
 
 
+def test_sdc_strip_format_jser():
+    instance = Config()
+    jser: JsonSerializer = JsonSerializer(write_only_format=True)
+    serialized = jser.json_serialize(instance)
+    assert serialized == {
+        "name": "default_name",
+        "save_model": True,
+        "batch_size": 64,
+        "__write_format__": "Config(SerializableDataclass)",
+    }
+    recovered = Config.load(serialized)
+    assert recovered == instance
+
+
 TYPE_MAP: dict[str, type] = {x.__name__: x for x in [int, float, str, bool]}