refactor migration

Author: weinbe58

src/bloqade/analog/migrate.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import IO, Any, Dict, List
 from dataclasses import field, dataclass
 
 import simplejson as json
@@ -26,7 +26,7 @@ def walk(self, obj: Dict[str, Any] | List[Any]) -> Dict[str, Any] | List[Any]:
         if isinstance(obj, dict):
             return self.walk_dict(obj)
         elif isinstance(obj, list):
-            return [self.walk(item) for item in obj]
+            return list(map(self.walk, obj))
         else:
             return obj
 
@@ -36,23 +36,22 @@ def convert(self, obj: Dict[str, Any] | List[Any]):
         return new_obj, self.has_done_something
 
 
+def _migrate(input_oi: IO[str], output_oi: IO[str], indent: int | None = None):
+    obj = json.load(input_oi)
+    new_obj, has_done_something = JSONWalker().convert(obj)
+
+    if has_done_something:
+        json.dump(new_obj, output_oi, indent=indent)
+
+
 def migrate(
     filename: str,
     indent: int | None = None,
     overwrite: bool = False,
 ):
-
-    with open(filename, "r") as io:
-        obj = json.load(io)
-        walker = JSONWalker()
-        new_obj, has_done_something = walker.convert(obj)
-
-    if has_done_something:
-        new_filename = (
-            filename if overwrite else filename.replace(".json", "-analog.json")
-        )
-        with open(new_filename, "w") as io:
-            json.dump(new_obj, io, indent=indent)
+    new_filename = filename if overwrite else filename.replace(".json", "-analog.json")
+    with open(filename, "r") as in_io, open(new_filename, "w") as out_io:
+        _migrate(in_io, out_io, indent)
 
 
 def _entry():

tests/test_migrate.py

@@ -1,8 +1,12 @@
-from bloqade.analog.migrate import JSONWalker
+import io
+
+import simplejson as json
+
+from bloqade.analog.migrate import _migrate
 
 
 def test_walk_dict():
-    walker = JSONWalker()
+
     obj = {
         "key1": "value1",
         "bloqade.key2": "value2",
@@ -19,5 +23,12 @@ def test_walk_dict():
         "list": [{"key6": "value6"}, {"bloqade.analog.key7": "value7"}],
     }
 
-    result = walker.walk_dict(obj)
-    assert result == expected
+    obj_str = json.dumps(obj)
+    expected_str = json.dumps(expected)
+
+    in_io = io.StringIO(obj_str)
+    out_io = io.StringIO()
+
+    _migrate(in_io, out_io)
+
+    assert out_io.getvalue() == expected_str