Use DataSource enum to validate source_id in the facade functions (#80)

* added validator for source_id and id

* validator for source id

* fix typehint error

* fix the failing test

Author: akash5100

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:
         """

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:
         """

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:
         """

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)
 
 

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,

hvpy/api_groups/jpeg2000/tests/test_get_jpx_closest_to_mid_point.py

@@ -2,32 +2,53 @@
 
 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://")
 
 
 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://")

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:
         """

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)
 
 

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,

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.