def test_collect_env_info(self):
info = collect_environment()
self.assertTrue(info["python"])
self.assertTrue(info["platform"])
self.assertGreater(len(info["packages"]), 0)
pkgs = set(p[0] for p in info["packages"])
self.assertIn("modelforge", pkgs)
self.assertIn("asdf", pkgs)
self.assertIn("numpy", pkgs)
def _write_tree(self,
tree: dict,
output: Union[str, BinaryIO],
file_mode: int = 0o666) -> None:
"""
Write the model to disk.
:param tree: The data dict - will be the ASDF tree.
:param output: The output file path or a file object.
:param file_mode: The output file's permissions.
:return: None
"""
self.meta["created_at"] = get_datetime_now()
meta = self.meta.copy()
meta["environment"] = collect_environment()
final_tree = {}
final_tree.update(tree)
final_tree["meta"] = meta
isfileobj = not isinstance(output, str)
if not isfileobj:
self._source = output
path = output
output = open(output, "wb")
os.chmod(path, file_mode)
pos = 0
else:
pos = output.tell()
try:
with asdf.AsdfFile(final_tree) as file:
queue = [("", tree)]
while queue:
path, element = queue.pop()
if isinstance(element, dict):
for key, val in element.items():
queue.append((path + "/" + key, val))
elif isinstance(element, (list, tuple)):
for child in element:
queue.append((path, child))
elif isinstance(element, numpy.ndarray):
path += "/"
if path not in self._compression_prefixes:
self._log.debug("%s -> %s compression", path,
self.ARRAY_COMPRESSION)
file.set_array_compression(element,
self.ARRAY_COMPRESSION)
else:
self._log.debug("%s -> compression disabled", path)
file.write_to(output)
self._size = output.seek(0, os.SEEK_END) - pos
finally:
if not isfileobj:
output.close()