SLEAP: Convert `dict` state to `tuple` state skeleton.json

sleap_skeleton_convert_dict_state_to_tuple_state.py

"""Converts the "py/state" entries in a JSON file to contain "py/tuple" entries.

SLEAP uses jsonpickle to encode/decode `Skeleton` objects to/from JSON. When a Python 
object is encoded to JSON, the "py/state" key is used to store the object's state. 
jsonpickle relies on the object's `__getstate__` method to determine the state to store.
attr creates a `__getstate__` method for the object. However, in `attr>=22.2.0`, the 
formatting changes s.t. `__getstate__` returns a dictionary of the object's attributes
instead of a tuple. This causes problems when decoding the JSON back to a Python object
if the system has an older version of `attr<22.2.0`.

This script converts the "py/state" entries in a JSON file to contain "py/tuple" entries
so that these skeleton.json files (created with `attr>=22.2.0`) can be read in with 
`attr<22.2.0`.

Example of input vs. output JSON:
Input JSON:
    {
        "py/state": {
            "name": "skeleton",
            "weight": 1.0
        }
    }

Output JSON:
    {
        "py/state": {
            "py/tuple": ["skeleton", 1.0]
        }
    }

Loading the JSON file with `attr<22.2.0`:
    In [1]: from sleap import Skeleton

    In [2]: input_json = "dict_state.json"

    In [3]: output_json = "tuple_state.json"

    In [4]: skeleton_input = Skeleton.load_json(input_json)

    In [5]: skeleton_input
    Out[5]: Skeleton(name='Skeleton-3', description='None', nodes=['name', 'name', 'name', 
        'name', 'name'], edges=[('name', 'name'), ('name', 'name'), ('name', 'name'), (
        'name', 'name'), ('name', 'name')], symmetries=[])

    In [6]: skeleton_output = Skeleton.load_json(output_json)

    In [7]: skeleton_output
    Out[7]: Skeleton(name='Skeleton-3', description='None', nodes=['nose', 'right-ear', 
        'left-ear', 'thorax', 'tail-base'], edges=[('nose', 'right-ear'), ('nose', 
        'left-ear'), ('right-ear', 'thorax'), ('left-ear', 'thorax'), ('thorax', 
        'tail-base')], symmetries=[])
"""

import json

def main(input_json: str, output_json: str):
    def rewrite_py_state(data):
        if isinstance(data, dict):
            for key, value in data.items():
                if key == "py/state" and isinstance(value, dict):
                    if "name" in value and "weight" in value:
                        data[key] = {"py/tuple": [value["name"], value["weight"]]}
                else:
                    rewrite_py_state(value)
        elif isinstance(data, list):
            for item in data:
                rewrite_py_state(item)

    # Read the JSON file
    with open(input_json, 'r') as file:
        data = json.load(file)

    # Rewrite the "py/state" entries
    rewrite_py_state(data)

    # Write the modified data back to the JSON file
    with open(output_json, 'w') as file:
        json.dump(data, file, indent=4)

if __name__ == '__main__':
    input_json = "dict_state.json"
    output_json = "tuple_state.json"
    main(input_json=input_json, output_json=output_json)