From 52e53f9e46adcb80e430561d8f489f9fcd93075d Mon Sep 17 00:00:00 2001 From: Matthew Iannucci Date: Mon, 21 Oct 2024 09:34:22 -0400 Subject: [PATCH] Some type stuff --- virtualizarr/tests/test_writers/test_icechunk.py | 3 ++- virtualizarr/writers/icechunk.py | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/virtualizarr/tests/test_writers/test_icechunk.py b/virtualizarr/tests/test_writers/test_icechunk.py index 596676bf..fd0284f0 100644 --- a/virtualizarr/tests/test_writers/test_icechunk.py +++ b/virtualizarr/tests/test_writers/test_icechunk.py @@ -149,6 +149,7 @@ def test_set_single_virtual_ref_with_encoding( root_group = group(store=icechunk_filestore) air_array = root_group["air"] + assert isinstance(air_array, Array) # check array metadata assert air_array.shape == (2920, 25, 53) @@ -159,7 +160,7 @@ def test_set_single_virtual_ref_with_encoding( # xarray performs this when cf_decoding is True, but we are not loading # with xarray here so we scale it manually. scale_factor = air_array.attrs["scale_factor"] - scaled_air_array = air_array[:] * scale_factor + scaled_air_array = air_array[:] * scale_factor # type: ignore # check chunk references # TODO we can't explicitly check that the path/offset/length is correct because icechunk doesn't yet expose any get_virtual_refs method diff --git a/virtualizarr/writers/icechunk.py b/virtualizarr/writers/icechunk.py index 6a361406..24bf7585 100644 --- a/virtualizarr/writers/icechunk.py +++ b/virtualizarr/writers/icechunk.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import numpy as np from xarray import Dataset @@ -101,7 +101,7 @@ def write_virtual_variable_to_icechunk( var: Variable, ) -> None: """Write a single virtual variable into an icechunk store""" - ma = var.data + ma = cast(ManifestArray, var.data) zarray = ma.zarray # creates array if it doesn't already exist @@ -154,7 +154,7 @@ def write_manifest_virtual_refs( "multi_index", "c_index", # TODO is "c_index" correct? what's the convention for zarr chunk keys? ], - op_flags=[["readonly"]] * 3, + op_flags=[["readonly"]] * 3, # type: ignore ) for path, offset, length in it: index = it.multi_index