sqme: mypy redo

opendatacube · Jun 24, 2024 · 32a75b5 · 32a75b5
1 parent 3cb07d3
commit 32a75b5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/odc/loader/test_memreader.py b/odc/loader/test_memreader.py
@@ -10,6 +10,7 @@
 import pytest
 import xarray as xr
 from odc.geo.data import country_geom
+from odc.geo.gcp import GCPGeoBox
 from odc.geo.geobox import GeoBox
 from odc.geo.xr import ODCExtensionDa, ODCExtensionDs, rasterize
 
@@ -37,6 +38,8 @@ def test_mem_reader(sample_ds: xr.Dataset) -> None:
     assert isinstance(sample_ds.odc, ODCExtensionDs)
     gbox = sample_ds.odc.geobox
     assert gbox is not None
+    assert isinstance(gbox, GeoBox)
+
     driver = XrMemReaderDriver(sample_ds)
 
     assert driver.md_parser is not None
@@ -173,6 +176,9 @@ def test_memreader_zarr(sample_ds: xr.Dataset):
     gbox = sample_ds.odc.geobox
     chunks = None
 
+    if isinstance(gbox, GCPGeoBox):
+        gbox = gbox.a
+
     md_store: dict[str, bytes] = {}
     chunk_store: dict[str, bytes] = {}
     sample_ds.to_zarr(md_store, chunk_store, compute=False, consolidated=True)

diff --git a/odc/loader/testing/mem_reader.py b/odc/loader/testing/mem_reader.py
@@ -11,7 +11,7 @@
 import fsspec
 import numpy as np
 import xarray as xr
-from odc.geo.geobox import GeoBox, GeoBoxBase
+from odc.geo.geobox import GeoBox
 from odc.geo.xr import ODCExtensionDa, ODCExtensionDs, xr_coords, xr_reproject
 
 from ..types import (
@@ -148,7 +148,7 @@ class Context:
 
     def __init__(
         self,
-        geobox: None | GeoBoxBase,
+        geobox: GeoBox,
         chunks: None | dict[str, int],
     ) -> None:
         self.geobox = geobox