Add createScreenshot function (#77)

* Add createScreenshot function

* Suggested changes for the filename and testing

* Apply suggestions from code review

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

* Applying suggested changes

* Applying suggested changes

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

Author: akash5100

hvpy/__init__.py

@@ -2,6 +2,6 @@
 from .datasource import *
 from .event import *
 from .facade import *
-from .helpers import createMovie
+from .helpers import createMovie, createScreenshot
 from .utils import create_events, create_layers, save_file
 from .version import __version__

hvpy/helpers.py

@@ -4,11 +4,13 @@
 from datetime import datetime
 
 from hvpy.api_groups.movies.queue_movie import queueMovieInputParameters
-from hvpy.facade import downloadMovie, getMovieStatus, queueMovie
+from hvpy.api_groups.screenshots.take_screenshot import takeScreenshotInputParameters
+from hvpy.facade import downloadMovie, downloadScreenshot, getMovieStatus, queueMovie, takeScreenshot
 from hvpy.utils import _add_shared_docstring, save_file
 
 __all__ = [
     "createMovie",
+    "createScreenshot",
 ]
 
 
@@ -57,7 +59,7 @@ def createMovie(
         Default is `False`.
     filename
         The path to save the file to.
-        Optional, will default to ``f"{starttime}_{endtime}.{format}"``.
+        Optional, will default to ``f"{res['id']}_{startTime.date()}_{endTime.date()}.{format}"``.
     hq
         Download a higher-quality movie file (valid for "mp4" movies only, ignored otherwise).
         Default is `False`, optional.
@@ -124,3 +126,84 @@ def createMovie(
         overwrite=overwrite,
     )
     return Path(filename)
+
+
+@_add_shared_docstring(takeScreenshotInputParameters)
+def createScreenshot(
+    date: datetime,
+    imageScale: float,
+    layers: str,
+    events: Optional[str] = None,
+    eventLabels: bool = False,
+    scale: bool = False,
+    scaleType: Optional[str] = None,
+    scaleX: Optional[int] = None,
+    scaleY: Optional[int] = None,
+    width: Optional[str] = None,
+    height: Optional[str] = None,
+    x0: Optional[str] = None,
+    y0: Optional[str] = None,
+    x1: Optional[str] = None,
+    y1: Optional[str] = None,
+    x2: Optional[str] = None,
+    y2: Optional[str] = None,
+    watermark: bool = False,
+    overwrite: bool = False,
+    filename: Optional[Union[str, Path]] = None,
+) -> Path:
+    """
+    Automatically creates a screenshot using `takeScreenshot`,
+    `downloadScreenshot` functions.
+
+    Parameters
+    ----------
+    overwrite
+        Whether to overwrite the file if it already exists.
+        Default is `False`.
+    filename
+        The path to save the file to.
+        Optional, will default to ``f"{res['id']}_{date.date()}.png"``.
+    {Insert}
+
+    Examples
+    --------
+    >>> from hvpy import createScreenshot, DataSource, create_events, create_layers, EventType
+    >>> from datetime import datetime, timedelta
+    >>> screenshot_location = createScreenshot(
+    ...     date=datetime.today() - timedelta(days=15),
+    ...     layers=create_layers([(DataSource.AIA_171, 100)]),
+    ...     events=create_events([EventType.ACTIVE_REGION]),
+    ...     eventLabels=True,
+    ...     imageScale=1,
+    ...     x0=0,
+    ...     y0=0,
+    ...     width=100,
+    ...     height=100,
+    ...     filename="my_screenshot",
+    ... )
+    >>> # This is to cleanup the file created from the example
+    >>> # you don't need to do this
+    >>> from pathlib import Path
+    >>> Path('my_screenshot.png').unlink()
+    """
+    input_params = locals()
+    # These are used later on but we want to avoid passing
+    # them into takeScreenshot.
+    input_params.pop("overwrite")
+    input_params.pop("filename")
+    res = takeScreenshot(**input_params)
+    if res.get("error"):
+        raise RuntimeError(res["error"])
+    binary_data = downloadScreenshot(
+        id=res["id"],
+    )
+    if filename is None:
+        filename = f"{res['id']}_{date.date()}.png"
+    else:
+        filename = f"{filename}.png"
+    save_file(
+        data=binary_data,
+        filename=filename,
+        overwrite=overwrite,
+    )
+    return Path(filename)

hvpy/tests/test_helper.py

@@ -4,7 +4,7 @@
 import pytest
 
 from hvpy.datasource import DataSource
-from hvpy.helpers import createMovie
+from hvpy.helpers import createMovie, createScreenshot
 from hvpy.utils import create_events, create_layers
 
 
@@ -93,3 +93,37 @@ def test_error_handling(tmp_path):
             imageScale=1,
             filename=f1,
         )
+
+
+def test_createScreenshot(date, tmp_path):
+    f1 = tmp_path / "screenshot"
+    result = createScreenshot(
+        date=date,
+        layers=create_layers([(DataSource.AIA_171, 100)]),
+        events=create_events(["AR"]),
+        imageScale=2.5,
+        x0=0,
+        y0=0,
+        width=4000,
+        height=4000,
+        filename=f1,
+    )
+    assert isinstance(result, Path)
+    assert result.exists()
+    assert result == tmp_path / "screenshot.png"
+
+
+def test_createScreenshot_with_none_filename(date):
+    result = createScreenshot(
+        date=date,
+        layers=create_layers([(DataSource.AIA_171, 100)]),
+        events=create_events(["AR"]),
+        imageScale=2.5,
+        x0=0,
+        y0=0,
+        width=4000,
+        height=4000,
+    )
+    assert isinstance(result, Path)
+    assert result.exists()
+    result.unlink()  # clean up