diff --git a/hvpy/__init__.py b/hvpy/__init__.py index a86fdde..14c3a62 100644 --- a/hvpy/__init__.py +++ b/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__ diff --git a/hvpy/helpers.py b/hvpy/helpers.py index a9e0136..0853c40 100644 --- a/hvpy/helpers.py +++ b/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) diff --git a/hvpy/tests/test_helper.py b/hvpy/tests/test_helper.py index 96a1188..9e53532 100644 --- a/hvpy/tests/test_helper.py +++ b/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