diff --git a/hvpy/api_groups/jpeg2000/get_jp2_image.py b/hvpy/api_groups/jpeg2000/get_jp2_image.py index a812643..af6b4e2 100644 --- a/hvpy/api_groups/jpeg2000/get_jp2_image.py +++ b/hvpy/api_groups/jpeg2000/get_jp2_image.py @@ -1,9 +1,11 @@ +from typing import Union from datetime import datetime from pydantic import Field, validator +from hvpy.datasource import DataSource from hvpy.io import HvpyParameters, OutputType -from hvpy.utils import convert_date_to_isoformat +from hvpy.utils import _data_source_to_int, convert_date_to_isoformat class getJP2ImageInputParameters(HvpyParameters): @@ -32,10 +34,11 @@ class getJP2ImageInputParameters(HvpyParameters): """ date: datetime - sourceId: int + sourceId: Union[int, DataSource] jpip: bool = False Json: bool = Field(False, alias="json") _date_validator = validator("date", allow_reuse=True)(convert_date_to_isoformat) + _source_id_validator = validator("sourceId", allow_reuse=True)(_data_source_to_int) def get_output_type(self) -> OutputType: """ diff --git a/hvpy/api_groups/jpeg2000/get_jpx.py b/hvpy/api_groups/jpeg2000/get_jpx.py index 01782c4..db59370 100644 --- a/hvpy/api_groups/jpeg2000/get_jpx.py +++ b/hvpy/api_groups/jpeg2000/get_jpx.py @@ -1,10 +1,11 @@ -from typing import Optional +from typing import Union, Optional from datetime import datetime from pydantic import validator +from hvpy.datasource import DataSource from hvpy.io import HvpyParameters, OutputType -from hvpy.utils import convert_date_to_isoformat +from hvpy.utils import _data_source_to_int, convert_date_to_isoformat class getJPXInputParameters(HvpyParameters): @@ -44,12 +45,13 @@ class getJPXInputParameters(HvpyParameters): startTime: datetime endTime: datetime - sourceId: int + sourceId: Union[int, DataSource] linked: bool = True verbose: bool = False jpip: bool = False cadence: Optional[int] = None _date_validator = validator("startTime", "endTime", allow_reuse=True)(convert_date_to_isoformat) + _source_id_validator = validator("sourceId", allow_reuse=True)(_data_source_to_int) def get_output_type(self) -> OutputType: """ diff --git a/hvpy/api_groups/jpeg2000/get_jpx_closest_to_mid_point.py b/hvpy/api_groups/jpeg2000/get_jpx_closest_to_mid_point.py index 4071614..d0e17aa 100644 --- a/hvpy/api_groups/jpeg2000/get_jpx_closest_to_mid_point.py +++ b/hvpy/api_groups/jpeg2000/get_jpx_closest_to_mid_point.py @@ -1,10 +1,11 @@ -from typing import List +from typing import List, Union from datetime import datetime from pydantic import validator +from hvpy.datasource import DataSource from hvpy.io import HvpyParameters, OutputType -from hvpy.utils import convert_date_to_unix +from hvpy.utils import _data_source_to_int, convert_date_to_unix class getJPXClosestToMidPointInputParameters(HvpyParameters): @@ -39,11 +40,12 @@ class getJPXClosestToMidPointInputParameters(HvpyParameters): startTimes: List[datetime] endTimes: List[datetime] - sourceId: int + sourceId: Union[int, DataSource] linked: bool = True verbose: bool = False jpip: bool = False _date_validator = validator("startTimes", "endTimes", allow_reuse=True)(convert_date_to_unix) + _source_id_validator = validator("sourceId", allow_reuse=True)(_data_source_to_int) def get_output_type(self) -> OutputType: """ diff --git a/hvpy/api_groups/jpeg2000/tests/test_get_jp2_image.py b/hvpy/api_groups/jpeg2000/tests/test_get_jp2_image.py index 855f800..ba53e64 100644 --- a/hvpy/api_groups/jpeg2000/tests/test_get_jp2_image.py +++ b/hvpy/api_groups/jpeg2000/tests/test_get_jp2_image.py @@ -2,10 +2,11 @@ from hvpy import getJP2Image from hvpy.api_groups.jpeg2000.get_jp2_image import getJP2ImageInputParameters +from hvpy.datasource import DataSource def test_str_response(date): - response = getJP2Image(date=date, sourceId=14, jpip=True, json=False) + response = getJP2Image(date=date, sourceId=DataSource.AIA_335, jpip=True, json=False) assert isinstance(response, str) assert response.startswith("jpip://") @@ -18,7 +19,7 @@ def test_json_response(date): def test_raw_response(date): - response = getJP2Image(date=date, sourceId=14, jpip=False, json=False) + response = getJP2Image(date=date, sourceId=DataSource.AIA_335, jpip=False, json=False) assert isinstance(response, bytes) @@ -28,7 +29,7 @@ def test_raw_response_with_json(date): def test_default_response(date): - response = getJP2Image(date=date, sourceId=14) + response = getJP2Image(date=date, sourceId=DataSource.AIA_335) assert isinstance(response, bytes) diff --git a/hvpy/api_groups/jpeg2000/tests/test_get_jpx.py b/hvpy/api_groups/jpeg2000/tests/test_get_jpx.py index 80fe3d5..4fa3c96 100644 --- a/hvpy/api_groups/jpeg2000/tests/test_get_jpx.py +++ b/hvpy/api_groups/jpeg2000/tests/test_get_jpx.py @@ -2,13 +2,14 @@ from hvpy import getJPX from hvpy.api_groups.jpeg2000.get_jpx import getJPXInputParameters +from hvpy.datasource import DataSource def test_raw_response(start_time, end_time): response = getJPX( startTime=start_time, endTime=end_time, - sourceId=14, + sourceId=DataSource.AIA_335, linked=False, verbose=False, jpip=False, @@ -21,7 +22,7 @@ def test_str_response(start_time, end_time): response = getJPX( startTime=start_time, endTime=end_time, - sourceId=14, + sourceId=DataSource.AIA_335, linked=False, verbose=False, jpip=True, diff --git a/hvpy/api_groups/jpeg2000/tests/test_get_jpx_closest_to_mid_point.py b/hvpy/api_groups/jpeg2000/tests/test_get_jpx_closest_to_mid_point.py index a97224b..55d9293 100644 --- a/hvpy/api_groups/jpeg2000/tests/test_get_jpx_closest_to_mid_point.py +++ b/hvpy/api_groups/jpeg2000/tests/test_get_jpx_closest_to_mid_point.py @@ -2,18 +2,29 @@ from hvpy import getJPXClosestToMidPoint from hvpy.api_groups.jpeg2000.get_jpx_closest_to_mid_point import getJPXClosestToMidPointInputParameters +from hvpy.datasource import DataSource def test_raw_response(start_times, end_times): response = getJPXClosestToMidPoint( - startTimes=start_times, endTimes=end_times, sourceId=14, linked=False, verbose=False, jpip=False + startTimes=start_times, + endTimes=end_times, + sourceId=DataSource.AIA_335, + linked=False, + verbose=False, + jpip=False, ) assert isinstance(response, bytes) def test_str_response(start_times, end_times): response = getJPXClosestToMidPoint( - startTimes=start_times, endTimes=end_times, sourceId=14, linked=False, verbose=False, jpip=True + startTimes=start_times, + endTimes=end_times, + sourceId=14, + linked=False, + verbose=False, + jpip=True, ) assert isinstance(response, str) assert response.startswith("jpip://") @@ -21,13 +32,23 @@ def test_str_response(start_times, end_times): def test_json_response(start_times, end_times): response = getJPXClosestToMidPoint( - startTimes=start_times, endTimes=end_times, sourceId=14, linked=False, verbose=True, jpip=True + startTimes=start_times, + endTimes=end_times, + sourceId=DataSource.AIA_335, + linked=False, + verbose=True, + jpip=True, ) assert isinstance(response, dict) assert response["uri"].startswith("jpip://") response = getJPXClosestToMidPoint( - startTimes=start_times, endTimes=end_times, sourceId=14, linked=False, verbose=True, jpip=False + startTimes=start_times, + endTimes=end_times, + sourceId=14, + linked=False, + verbose=True, + jpip=False, ) assert isinstance(response, dict) assert response["uri"].startswith("https://") diff --git a/hvpy/api_groups/official_clients/get_closest_image.py b/hvpy/api_groups/official_clients/get_closest_image.py index 10bfe65..8df873f 100644 --- a/hvpy/api_groups/official_clients/get_closest_image.py +++ b/hvpy/api_groups/official_clients/get_closest_image.py @@ -1,10 +1,11 @@ -from typing import Optional +from typing import Union, Optional from datetime import datetime from pydantic import validator +from hvpy.datasource import DataSource from hvpy.io import HvpyParameters, OutputType -from hvpy.utils import convert_date_to_isoformat +from hvpy.utils import _data_source_to_int, convert_date_to_isoformat class getClosestImageInputParameters(HvpyParameters): @@ -30,9 +31,10 @@ class getClosestImageInputParameters(HvpyParameters): """ date: datetime - sourceId: int + sourceId: Union[int, DataSource] callback: Optional[str] = None _date_validator = validator("date", allow_reuse=True)(convert_date_to_isoformat) + _source_id_validator = validator("sourceId", allow_reuse=True)(_data_source_to_int) def get_output_type(self) -> OutputType: """ diff --git a/hvpy/api_groups/official_clients/tests/test_get_closest_image.py b/hvpy/api_groups/official_clients/tests/test_get_closest_image.py index 5f659e6..c65739a 100644 --- a/hvpy/api_groups/official_clients/tests/test_get_closest_image.py +++ b/hvpy/api_groups/official_clients/tests/test_get_closest_image.py @@ -2,10 +2,11 @@ from hvpy import getClosestImage from hvpy.api_groups.official_clients.get_closest_image import getClosestImageInputParameters +from hvpy.datasource import DataSource def test_json_res(date): - response = getClosestImage(date=date, sourceId=14) + response = getClosestImage(date=date, sourceId=DataSource.AIA_335) assert isinstance(response, dict) diff --git a/hvpy/facade.py b/hvpy/facade.py index 6d6c5a7..0609ef9 100644 --- a/hvpy/facade.py +++ b/hvpy/facade.py @@ -2,6 +2,7 @@ from datetime import datetime from hvpy.core import execute_api_call +from hvpy.datasource import DataSource from hvpy.parameters import * from hvpy.utils import _add_shared_docstring @@ -28,7 +29,7 @@ @_add_shared_docstring(getJP2ImageInputParameters) def getJP2Image( date: datetime, - sourceId: int, + sourceId: Union[int, DataSource], jpip: bool = False, json: bool = False, ) -> Union[bytes, str, Dict[str, Any]]: @@ -75,7 +76,7 @@ def getJP2Header( def getJPXClosestToMidPoint( startTimes: List[datetime], endTimes: List[datetime], - sourceId: int, + sourceId: Union[int, DataSource], linked: bool = True, verbose: bool = False, jpip: bool = False, @@ -115,7 +116,7 @@ def getJPXClosestToMidPoint( def getJPX( startTime: List[datetime], endTime: List[datetime], - sourceId: int, + sourceId: Union[int, DataSource], linked: bool = True, verbose: bool = False, jpip: bool = False, @@ -175,7 +176,7 @@ def getStatus() -> Union[bytes, str, Dict[str, Any]]: @_add_shared_docstring(getClosestImageInputParameters) def getClosestImage( date: datetime, - sourceId: int, + sourceId: Union[int, DataSource], callback: Optional[str] = None, ) -> Union[bytes, str, Dict[str, Any]]: """ @@ -424,7 +425,7 @@ def reQueueMovie( @_add_shared_docstring(getMovieStatusInputParameters) def getMovieStatus( - id: str, + id: Union[int, DataSource], format: str, verbose: bool = False, callback: Optional[str] = None, @@ -454,7 +455,7 @@ def getMovieStatus( @_add_shared_docstring(downloadMovieInputParameters) def downloadMovie( - id: str, + id: Union[int, DataSource], format: str, hq: bool = False, ) -> Union[bytes, str, Dict[str, Any]]: @@ -527,7 +528,7 @@ def shortenURL( @_add_shared_docstring(getTileInputParameters) def getTile( - id: int, + id: Union[int, DataSource], x: int, y: int, imageScale: int, diff --git a/hvpy/utils.py b/hvpy/utils.py index 218fad7..484d914 100644 --- a/hvpy/utils.py +++ b/hvpy/utils.py @@ -49,6 +49,18 @@ def convert_date_to_unix(v: list) -> str: return ",".join([str(int(datetime.timestamp(d))) for d in v]) +def _data_source_to_int(source: Union[int, DataSource]) -> int: + """ + Converts a `~hvpy.DataSource` to an integer. + + Parameters + ---------- + source + The `~hvpy.DataSource` to convert. + """ + return _to_datasource(source).value + + def _to_datasource(val: Union[int, DataSource]) -> DataSource: """ Validates the input and converts it to a DataSource enum.