TYP: Add type hints for 'data_source' in load_earth_relief and load_e…

…arth_magnetic_anomaly functions (#2849)
GenericMappingTools · Dec 8, 2023 · 3812c99 · 3812c99
1 parent 7845ae8
commit 3812c99
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 18 deletions.
diff --git a/pygmt/datasets/earth_magnetic_anomaly.py b/pygmt/datasets/earth_magnetic_anomaly.py
@@ -4,6 +4,8 @@
 
 The grids are available in various resolutions.
 """
+from typing import Literal
+
 from pygmt.datasets.load_remote_dataset import _load_remote_dataset
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import kwargs_to_strings
@@ -13,7 +15,10 @@
 
 @kwargs_to_strings(region="sequence")
 def load_earth_magnetic_anomaly(
-    resolution="01d", region=None, registration=None, data_source="emag2"
+    resolution="01d",
+    region=None,
+    registration=None,
+    data_source: Literal["emag2", "emag2_4km", "wdmam"] = "emag2",
 ):
     r"""
     Load the Earth magnetic anomaly datasets in various resolutions.
@@ -77,17 +82,15 @@ def load_earth_magnetic_anomaly(
         for all resolutions except ``"02m"`` for ``data_source="emag2"`` or
         ``data_source="emag2_4km"``, which are ``"pixel"`` only.
 
-    data_source : str
+    data_source
         Select the source of the magnetic anomaly data. Available options are:
 
-        - ``"emag2"``: EMAG2 Earth Magnetic Anomaly Model [Default
-          option]. It only includes data observed at sea level over
-          oceanic regions. See :gmt-datasets:`earth-mag.html`.
-
+        - ``"emag2"``: EMAG2 Earth Magnetic Anomaly Model. It only includes
+          data observed at sea level over oceanic regions.
+          See :gmt-datasets:`earth-mag.html`.
         - ``"emag2_4km"``: Use a version of EMAG2 where all observations
           are relative to an altitude of 4 km above the geoid and include
           data over land.
-
         - ``"wdmam"``: World Digital Magnetic Anomaly Map (WDMAM).
           See :gmt-datasets:`earth-wdmam.html`.
 

diff --git a/pygmt/datasets/earth_relief.py b/pygmt/datasets/earth_relief.py
@@ -4,6 +4,8 @@
 
 The grids are available in various resolutions.
 """
+from typing import Literal
+
 from pygmt.datasets.load_remote_dataset import _load_remote_dataset
 from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import kwargs_to_strings
@@ -16,7 +18,7 @@ def load_earth_relief(
     resolution="01d",
     region=None,
     registration=None,
-    data_source="igpp",
+    data_source: Literal["igpp", "gebco", "gebcosi", "synbath"] = "igpp",
     use_srtm=False,
 ):
     r"""
@@ -74,22 +76,18 @@ def load_earth_relief(
         ``"gridline"`` for gridline registration. Default is ``"gridline"``
         for all resolutions except ``"15s"`` which is ``"pixel"`` only.
 
-    data_source : str
+    data_source
         Select the source for the Earth relief data. Available options are:
 
-        - ``"igpp"``: IGPP Earth Relief [Default option]. See
+        - ``"igpp"``: IGPP Earth Relief. See
           :gmt-datasets:`earth-relief.html`.
-
-        - ``"synbath"``: IGPP Earth Relief dataset that uses
-          stastical properties of young seafloor to provide a more realistic
-          relief of young areas with small seamounts.
-
+        - ``"synbath"``: IGPP Earth Relief dataset that uses stastical
+          properties of young seafloor to provide a more realistic relief
+          of young areas with small seamounts.
         - ``"gebco"``: GEBCO Earth Relief with only observed relief and
           inferred relief via altimetric gravity. See
           :gmt-datasets:`earth-gebco.html`.
-
-        - ``"gebcosi"``: GEBCO Earth Relief that gives sub-ice (si)
-          elevations.
+        - ``"gebcosi"``: GEBCO Earth Relief that gives sub-ice (si) elevations.
 
     use_srtm : bool
         By default, the land-only SRTM tiles from NASA are used to generate the