refactor(Workspace): default zipping or not based on filename

Author: yt-ms

src/structurizr/workspace.py

@@ -196,7 +196,7 @@ def __init__(
 
     @classmethod
     def load(cls, filename: Union[str, Path]) -> "Workspace":
-        """Load a workspace from a file (which may optionally be gzipped)."""
+        """Load a workspace from a JSON file (which may optionally be gzipped)."""
         filename = Path(filename)
         try:
             with gzip.open(filename, "rt") as handle:
@@ -217,20 +217,25 @@ def dump(
         self,
         filename: Union[str, Path],
         *,
-        zip: bool = False,
+        zip: Optional[bool] = None,
         indent: Optional[int] = None,
         **kwargs
     ):
         """
-        Save a workspace to a file, optionally zipped.
+        Save a workspace as JSON to a file, optionally gzipped.
+
+        By default, filenames ending with `.gz` will be zipped and anything else won't,
+        however this can be overridden by explicitly passing the `zip` argument.
 
         Arguments:
-            filename (str/Path): filename to write to.
-            zip (bool): if true then contents will be zipped with GZip.
-            indent (int): if specified then pretty-print the JSON with given indent.
+            filename: filename to write to.
+            zip: if specified then controls whether the contents are gzipped.
+            indent: if specified then pretty-print the JSON with given indent.
             kwargs: other arguments to pass through to `json.dumps()`.
         """
         filename = Path(filename)
+        if zip is None:
+            zip = str(filename).endswith(".gz")
         with gzip.open(filename, "wt") if zip else filename.open("wt") as handle:
             handle.write(self.dumps(indent=indent, **kwargs))
 

tests/integration/test_workspace_io.py

@@ -136,14 +136,30 @@ def test_save_and_load_workspace_to_gzipped_file(monkeypatch, tmp_path: Path):
 
     filepath = tmp_path / "test_workspace.json.gz"
 
-    workspace.dump(filepath, zip=True)
+    workspace.dump(filepath)
     workspace2 = Workspace.load(filepath)
 
     expected = WorkspaceIO.from_orm(workspace)
     actual = WorkspaceIO.from_orm(workspace2)
     assert json.loads(actual.json()) == json.loads(expected.json())
 
 
+def test_workspace_overridding_zip_flag(monkeypatch, tmp_path: Path):
+    """Test that default zipping can be overridden explicitly."""
+    monkeypatch.syspath_prepend(EXAMPLES)
+    example = import_module("getting_started")
+    workspace = example.main()
+
+    filepath = tmp_path / "test_workspace.json.gz"
+
+    workspace.dump(filepath, zip=False)
+    contents = filepath.read_text()
+    assert "My software system" in contents
+
+    # Make sure can be loaded even though its not zipped and ends with .gz
+    Workspace.load(filepath)
+
+
 def test_load_unknown_file_raises_file_not_found():
     """Test that attempting to load a non-existent file raises FileNotFound."""
     with pytest.raises(FileNotFoundError):