Save binary response with save_file function (#72)

* Save file function

* overwrite option for already existing function

* test with tmp_path

* update typehint for filename and tests

* Apply suggestions from code review

Co-authored-by: Nabil Freij <[email protected]>

* Fix mypy error

* test for string input

Co-authored-by: Nabil Freij <[email protected]>

Author: akash5100

hvpy/tests/test_utils.py

@@ -1,13 +1,16 @@
+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,
     _to_datasource,
     _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()

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)