diff --git a/hvpy/tests/test_utils.py b/hvpy/tests/test_utils.py index 78a335d..10a48b3 100644 --- a/hvpy/tests/test_utils.py +++ b/hvpy/tests/test_utils.py @@ -1,6 +1,8 @@ +from datetime import datetime + import pytest -from hvpy import DataSource, EventType +from hvpy import DataSource, EventType, takeScreenshot from hvpy.utils import ( _create_events_string, _create_layer_string, @@ -8,6 +10,7 @@ _to_event_type, create_events, create_layers, + save_file, ) @@ -65,3 +68,27 @@ def test_create_events(): def test_create_events_string(): assert _create_events_string(EventType.ACTIVE_REGION) == "[AR,all,1]" assert _create_events_string(EventType.ACTIVE_REGION, "xyz") == "[AR,xyz,1]" + + +def test_save_file(tmp_path): + f1 = tmp_path / "test.png" + res = takeScreenshot( + date=datetime.today(), + imageScale=2.44, + layers="[10,1,100]", + x0=0, + y0=0, + width=1920, + height=1200, + display=True, + ) + save_file(res, f1, overwrite=False) + assert f1.exists() + with pytest.raises(ValueError, match="already exists"): + save_file(res, f1, overwrite=False) + save_file(res, f1, overwrite=True) + assert f1.exists() + + f2 = tmp_path / "test2.png" + save_file(res, str(f2), overwrite=False) + assert f2.exists() diff --git a/hvpy/utils.py b/hvpy/utils.py index fcc2469..0c2a3d3 100644 --- a/hvpy/utils.py +++ b/hvpy/utils.py @@ -1,4 +1,5 @@ from typing import Any, List, Union, Callable +from pathlib import Path from datetime import datetime from hvpy.datasource import DataSource @@ -9,6 +10,7 @@ "convert_date_to_unix", "create_layers", "create_events", + "save_file", ] @@ -127,3 +129,24 @@ def create_events(events: List[Union[EventType, str, tuple]]) -> str: raise ValueError(f"{event} is not a EventType or str or two-length tuple") # Strips the final comma return constructed_events[:-1] + + +def save_file(data: bytearray, filename: Union[Path, str], overwrite: bool = False) -> None: + """ + Saves a file to the specified path. + + Parameters + ---------- + data + The data to save. + filename + The path to save the file to. + overwrite + Whether to overwrite the file if it already exists. + Default is `False`. + """ + if isinstance(filename, str): + filename = Path(filename) + if filename.exists() and not overwrite: + raise ValueError(f"{filename} already exists. Use overwrite=True to overwrite.") + filename.write_bytes(data)