From 170a15d71df17747029c1b37399c50d6a2a4fc9c Mon Sep 17 00:00:00 2001 From: Dongdong Tian Date: Tue, 17 Dec 2024 23:31:52 +0800 Subject: [PATCH] Use enums GridReg and GridType in test_accessor.py --- pygmt/tests/test_accessor.py | 45 ++++++++++++++++++------------------ 1 file changed, 23 insertions(+), 22 deletions(-) diff --git a/pygmt/tests/test_accessor.py b/pygmt/tests/test_accessor.py index 2701d2a0b3a..f7c3bac26ec 100644 --- a/pygmt/tests/test_accessor.py +++ b/pygmt/tests/test_accessor.py @@ -11,6 +11,7 @@ from pygmt import which from pygmt.clib import __gmt_version__ from pygmt.datasets import load_earth_relief +from pygmt.enums import GridReg, GridType from pygmt.exceptions import GMTInvalidInput @@ -21,8 +22,8 @@ def test_accessor_gridline_cartesian(): """ fname = which(fname="@test.dat.nc", download="a") grid = xr.open_dataarray(fname) - assert grid.gmt.registration == 0 # gridline registration - assert grid.gmt.gtype == 0 # cartesian coordinate type + assert grid.gmt.registration == GridReg.GRIDLINE + assert grid.gmt.gtype == GridType.CARTESIAN def test_accessor_pixel_geographic(): @@ -32,8 +33,8 @@ def test_accessor_pixel_geographic(): """ fname = which(fname="@earth_relief_01d_p", download="a") grid = xr.open_dataarray(fname, engine="netcdf4") - assert grid.gmt.registration == 1 # pixel registration - assert grid.gmt.gtype == 1 # geographic coordinate type + assert grid.gmt.registration == GridReg.PIXEL + assert grid.gmt.gtype == GridType.GEOGRAPHIC def test_accessor_set_pixel_registration(): @@ -41,9 +42,9 @@ def test_accessor_set_pixel_registration(): Check that we can set a grid to be Pixel registered with a registration value of 1. """ grid = xr.DataArray(data=[[0.1, 0.2], [0.3, 0.4]]) - assert grid.gmt.registration == 0 # default to gridline registration - grid.gmt.registration = 1 # set to pixel registration - assert grid.gmt.registration == 1 # ensure changed to pixel registration + assert grid.gmt.registration == GridReg.GRIDLINE + grid.gmt.registration = GridReg.PIXEL + assert grid.gmt.registration == GridReg.PIXEL @pytest.mark.benchmark @@ -53,11 +54,11 @@ def test_accessor_set_geographic_cartesian_roundtrip(): using a gtype of 1, set it to Geographic 0, and then back to Cartesian again 1. """ grid = xr.DataArray(data=[[0.1, 0.2], [0.3, 0.4]]) - assert grid.gmt.gtype == 0 # default to cartesian coordinate type - grid.gmt.gtype = 1 # set to geographic type - assert grid.gmt.gtype == 1 # ensure changed to geographic coordinate type - grid.gmt.gtype = 0 # set back to cartesian type - assert grid.gmt.gtype == 0 # ensure changed to cartesian coordinate type + assert grid.gmt.gtype == GridType.CARTESIAN + grid.gmt.gtype = GridType.GEOGRAPHIC + assert grid.gmt.gtype == GridType.GEOGRAPHIC + grid.gmt.gtype = GridType.CARTESIAN + assert grid.gmt.gtype == GridType.CARTESIAN def test_accessor_set_non_boolean(): @@ -93,8 +94,8 @@ def test_accessor_sliced_datacube(): with xr.open_dataset(fname) as dataset: grid = dataset.sel(level=500, month=1, drop=True).z - assert grid.gmt.registration == 0 # gridline registration - assert grid.gmt.gtype == 1 # geographic coordinate type + assert grid.gmt.registration == GridReg.GRIDLINE + assert grid.gmt.gtype == GridType.GEOGRAPHIC finally: Path(fname).unlink() @@ -109,19 +110,19 @@ def test_accessor_grid_source_file_not_exist(): resolution="05m", region=[0, 5, -5, 5], registration="pixel" ) # Registration and gtype are correct - assert grid.gmt.registration == 1 - assert grid.gmt.gtype == 1 + assert grid.gmt.registration == GridReg.PIXEL + assert grid.gmt.gtype == GridType.GEOGRAPHIC # The source grid file is undefined. assert grid.encoding.get("source") is None # For a sliced grid, fallback to default registration and gtype, # because the source grid file doesn't exist. sliced_grid = grid[1:3, 1:3] - assert sliced_grid.gmt.registration == 0 - assert sliced_grid.gmt.gtype == 0 + assert sliced_grid.gmt.registration == GridReg.GRIDLINE + assert sliced_grid.gmt.gtype == GridType.CARTESIAN # Still possible to manually set registration and gtype - sliced_grid.gmt.registration = 1 - sliced_grid.gmt.gtype = 1 - assert sliced_grid.gmt.registration == 1 - assert sliced_grid.gmt.gtype == 1 + sliced_grid.gmt.registration = GridReg.PIXEL + sliced_grid.gmt.gtype = GridType.GEOGRAPHIC + assert sliced_grid.gmt.registration == GridReg.PIXEL + assert sliced_grid.gmt.gtype == GridType.GEOGRAPHIC