def on_save(self, path=None): """ This is a hook that is called just before saving the file. It can be used, for example, to update values in the metadata that are based on the content of the data. Override it in the subclass to make it do something, but don't forget to "chain up" to the base class, since it does things there, too. Parameters ---------- path : str The path to the file that we're about to save to. """ if isinstance(path, str): self.meta.filename = os.path.basename(path) current_date = Time(datetime.datetime.now()) current_date.format = 'isot' self.meta.date = current_date.value # Enforce model_type to be the actual type of model being saved. self.meta.model_type = self._model_type # Remove None nodes from the tree: def _remove_none(node): if node is None: return RemoveNode else: return node self._instance = treeutil.walk_and_modify(self._instance, _remove_none)
def test_walk_and_modify_shared_references(): target = {"foo": "bar"} nested_in_dict = {"target": target} nested_in_list = [target] tree = { "target": target, "nested_in_dict": nested_in_dict, "nested_in_list": nested_in_list } assert tree["target"] is tree["nested_in_dict"]["target"] assert tree["target"] is tree["nested_in_list"][0] def _callback(node): if "foo" in node: return {"foo": "baz"} else: return node result = treeutil.walk_and_modify(tree, _callback) assert result is not tree assert result["target"] is not target assert result["target"]["foo"] == "baz" assert result["target"] is result["nested_in_dict"]["target"] assert result["target"] is result["nested_in_list"][0]
def test_walk_and_modify_retain_none(): tree = {'foo': 42, 'bar': None} def func(x): if x == 42: return None return x tree2 = treeutil.walk_and_modify(tree, func) assert tree2['foo'] is None assert tree2['bar'] is None
def test_walk_and_modify_remove_keys(): tree = {'foo': 42, 'bar': 43} def func(x): if x == 42: return treeutil.RemoveNode return x tree2 = treeutil.walk_and_modify(tree, func) assert 'foo' not in tree2 assert 'bar' in tree2
def test_walk_and_modify_remove_keys(): tree = { 'foo': 42, 'bar': 43 } def func(x): if x == 42: return None return x tree2 = treeutil.walk_and_modify(tree, func) assert 'foo' not in tree2 assert 'bar' in tree2
def _save_from_schema(hdulist, tree, schema): def convert_datetimes(node, json_id): if isinstance(node, datetime.datetime): node = time.Time(node) if isinstance(node, time.Time): node = str(time.Time(node, format='iso')) return node tree = treeutil.walk_and_modify(tree, convert_datetimes) validator = asdf_schema.get_validator( schema, None, FITS_VALIDATORS, FITS_SCHEMA_URL_MAPPING) validator.hdulist = hdulist # TODO: Handle comment stack on per-hdu-basis validator.comment_stack = [] # This actually kicks off the saving validator.validate(tree, _schema=schema)
def remove_none_from_tree(tree): """ Remove None values from a tree. Both dictionary keys and list indices with None values will be removed. Parameters ---------- tree : object The root node of the tree. Returns ------- object Modified tree. """ def _remove_none(node): if node is None: return RemoveNode else: return node return treeutil.walk_and_modify(tree, _remove_none)
from asdf.treeutil import walk_and_modify tree = { "a": { "aa": 1, "ab": 3 }, "b": 6, "c": { "ca": 5, "cb": 7, "cc": { "cca": 3, "ccb": 8, "ccc": 9 }, "d": 2 }, } def callback(node, json_id): if isinstance(node, int) and node == 3: print(f"json id: {json_id} / node: {node} -> -5") return -5 print(f"json id: {json_id} / node: {node}") return node walk_and_modify(tree, callback, postorder=False)