Skip to content

Commit

Permalink
feat(system): Added new method to save the system to a given folder.
Browse files Browse the repository at this point in the history
List of changes:
- Created a new method `.save` that saves all the content of the system
  (.json + time series) to a desired folder,
- Save method can archive the folder to enable easy sharing of the
  system,
- Added testing to validate that save in fact creates the folder and
  that when we zip it we delete that folder.

Closes #10
  • Loading branch information
pesap committed May 1, 2024
1 parent fe2295d commit d2e52c8
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 0 deletions.
53 changes: 53 additions & 0 deletions src/infrasys/system.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""Defines a System"""

import json
import shutil
from operator import itemgetter
from collections import defaultdict
from datetime import datetime
Expand Down Expand Up @@ -284,6 +285,58 @@ def from_dict(
logger.info("Deserialized system {}", system.label)
return system

def save(
self,
fpath: Path | str,
filename: str = "system.json",
zip: bool = False,
overwrite: bool = False,
) -> None:
"""Save the contents of a system and the Time series in a single directory.
By default, this method creates the user specified folder using the
`to_json` method. If user sets `zip = True`, we create the folder of
the user (if it does not exists), zip it to the same location specified
and delete the folder.
Parameters
----------
fpath : Path | str
Filepath to write the contents of the system.
zip : bool
Set to True if you want to archive to a zip file.
filename: str
Name of the sytem to serialize. Default value: "system.json".
overwrite: bool
Overwrites the system if it already exist on the fpath.
Examples
--------
>>> system.save("/home/$USER/systems/my_system")
INFO: Wrote system data to systems/system.json
INFO: Copied time series data to systems/my_system/system_time_series
>>> system.save("/home/$USER/systems/my_system", zip=True)
INFO: Wrote system data to systems/system1.json
INFO: Copied time series data to systems/system1_time_series
See Also
--------
to_json: System serialization
"""
if isinstance(fpath, str):
fpath = Path(fpath)
self.to_json(fpath / filename, overwrite=overwrite)

if zip:
logger.info("Archiving system and time series into a single file")
_ = shutil.make_archive(str(fpath), "zip", fpath)
logger.debug("Removing {}", fpath)
shutil.rmtree(fpath)

return

def add_component(self, component: Component, **kwargs) -> None:
"""Add one component to the system.
Expand Down
18 changes: 18 additions & 0 deletions tests/test_serialization.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import json
import random
import os
from datetime import datetime, timedelta

import numpy as np
Expand Down Expand Up @@ -187,3 +188,20 @@ def test_system_with_time_series_normalization(tmp_path, in_memory):
def test_json_schema():
schema = ComponentWithPintQuantity.model_json_schema()
assert isinstance(json.loads(json.dumps(schema)), dict)


def test_system_save(tmp_path, simple_system_with_time_series):
simple_system = simple_system_with_time_series
custom_folder = "my_system"
fpath = tmp_path / custom_folder
fname = "test_system"
simple_system.save(fpath, filename=fname)
assert os.path.exists(fpath), f"Folder {fpath} was not created successfully"
assert os.path.exists(fpath / fname), f"Serialized system {fname} was not created successfully"

custom_folder = "my_system_zip"
fpath = tmp_path / custom_folder
simple_system.save(fpath, filename=fname, zip=True)
assert not os.path.exists(fpath), f"Original folder {fpath} was not deleted sucessfully."
zip_fpath = f"{fpath}.zip"
assert os.path.exists(zip_fpath), f"Zip file {zip_fpath} does not exists"

0 comments on commit d2e52c8

Please sign in to comment.