55import sys
66from contextlib import suppress
77from textwrap import dedent
8+ from typing import Any
89
910import numpy as np
1011import pandas as pd
@@ -94,13 +95,14 @@ def test_copy(self):
9495 self .assertLazyAndIdentical (self .eager_var , self .lazy_var .copy (deep = True ))
9596
9697 def test_chunk (self ):
97- for chunks , expected in [
98+ test_cases : list [ tuple [ int | dict [ str , Any ], tuple [ tuple [ int , ...], ...]]] = [
9899 ({}, ((2 , 2 ), (2 , 2 , 2 ))),
99100 (3 , ((3 , 1 ), (3 , 3 ))),
100101 ({"x" : 3 , "y" : 3 }, ((3 , 1 ), (3 , 3 ))),
101102 ({"x" : 3 }, ((3 , 1 ), (2 , 2 , 2 ))),
102103 ({"x" : (3 , 1 )}, ((3 , 1 ), (2 , 2 , 2 ))),
103- ]:
104+ ]
105+ for chunks , expected in test_cases :
104106 rechunked = self .lazy_var .chunk (chunks )
105107 assert rechunked .chunks == expected
106108 self .assertLazyAndIdentical (self .eager_var , rechunked )
@@ -258,7 +260,7 @@ def test_missing_methods(self):
258260 with pytest .raises (NotImplementedError , match = "dask" ):
259261 v .argsort ()
260262 with pytest .raises (NotImplementedError , match = "dask" ):
261- v [0 ].item ()
263+ v [0 ].item () # type: ignore[attr-defined]
262264
263265 def test_univariate_ufunc (self ):
264266 u = self .eager_var
@@ -298,7 +300,7 @@ def test_persist(self):
298300
299301 (v2 ,) = dask .persist (v )
300302 assert v is not v2
301- assert len (v2 .__dask_graph__ ()) < len (v .__dask_graph__ ())
303+ assert len (v2 .__dask_graph__ ()) < len (v .__dask_graph__ ()) # type: ignore[arg-type]
302304 assert v2 .__dask_keys__ () == v .__dask_keys__ ()
303305 assert dask .is_dask_collection (v )
304306 assert dask .is_dask_collection (v2 )
@@ -345,7 +347,9 @@ def setUp(self):
345347 )
346348
347349 def test_chunk (self ) -> None :
348- for chunks , expected in [
350+ test_cases : list [
351+ tuple [int | str | dict [str , Any ], tuple [tuple [int , ...], ...]]
352+ ] = [
349353 ({}, ((2 , 2 ), (2 , 2 , 2 ))),
350354 (3 , ((3 , 1 ), (3 , 3 ))),
351355 ({"x" : 3 , "y" : 3 }, ((3 , 1 ), (3 , 3 ))),
@@ -354,7 +358,8 @@ def test_chunk(self) -> None:
354358 ({"x" : "16B" }, ((1 , 1 , 1 , 1 ), (2 , 2 , 2 ))),
355359 ("16B" , ((1 , 1 , 1 , 1 ), (1 ,) * 6 )),
356360 ("16MB" , ((4 ,), (6 ,))),
357- ]:
361+ ]
362+ for chunks , expected in test_cases :
358363 # Test DataArray
359364 rechunked = self .lazy_array .chunk (chunks )
360365 assert rechunked .chunks == expected
@@ -367,7 +372,7 @@ def test_chunk(self) -> None:
367372 lazy_dataset = self .lazy_array .to_dataset ()
368373 eager_dataset = self .eager_array .to_dataset ()
369374 expected_chunksizes = dict (zip (lazy_dataset .dims , expected , strict = True ))
370- rechunked = lazy_dataset .chunk (chunks )
375+ rechunked = lazy_dataset .chunk (chunks ) # type: ignore[assignment]
371376
372377 # Dataset.chunks has a different return type to DataArray.chunks - see issue #5843
373378 assert rechunked .chunks == expected_chunksizes
@@ -601,11 +606,12 @@ def test_reindex(self):
601606 u = self .eager_array .assign_coords (y = range (6 ))
602607 v = self .lazy_array .assign_coords (y = range (6 ))
603608
604- for kwargs in [
609+ kwargs_list : list [ dict [ str , Any ]] = [
605610 {"x" : [2 , 3 , 4 ]},
606611 {"x" : [1 , 100 , 2 , 101 , 3 ]},
607612 {"x" : [2.5 , 3 , 3.5 ], "y" : [2 , 2.5 , 3 ]},
608- ]:
613+ ]
614+ for kwargs in kwargs_list :
609615 expected = u .reindex (** kwargs )
610616 actual = v .reindex (** kwargs )
611617 self .assertLazyAndAllClose (expected , actual )
@@ -666,7 +672,9 @@ def test_stack(self):
666672 data = da .random .normal (size = (2 , 3 , 4 ), chunks = (1 , 3 , 4 ))
667673 arr = DataArray (data , dims = ("w" , "x" , "y" ))
668674 stacked = arr .stack (z = ("x" , "y" ))
669- z = pd .MultiIndex .from_product ([np .arange (3 ), np .arange (4 )], names = ["x" , "y" ])
675+ z = pd .MultiIndex .from_product (
676+ [list (range (3 )), list (range (4 ))], names = ["x" , "y" ]
677+ )
670678 expected = DataArray (data .reshape (2 , - 1 ), {"z" : z }, dims = ["w" , "z" ])
671679 assert stacked .data .chunks == expected .data .chunks
672680 self .assertLazyAndEqual (expected , stacked )
@@ -1167,10 +1175,10 @@ def returns_numpy(darray):
11671175 xr .map_blocks (returns_numpy , map_da )
11681176
11691177 with pytest .raises (TypeError , match = r"args must be" ):
1170- xr .map_blocks (operator .add , map_da , args = 10 )
1178+ xr .map_blocks (operator .add , map_da , args = 10 ) # type: ignore[arg-type]
11711179
11721180 with pytest .raises (TypeError , match = r"kwargs must be" ):
1173- xr .map_blocks (operator .add , map_da , args = [10 ], kwargs = [20 ])
1181+ xr .map_blocks (operator .add , map_da , args = [10 ], kwargs = [20 ]) # type: ignore[arg-type]
11741182
11751183 def really_bad_func (darray ):
11761184 raise ValueError ("couldn't do anything." )
@@ -1442,7 +1450,7 @@ def test_map_blocks_errors_bad_template(obj):
14421450 with pytest .raises (ValueError , match = r"Received dimension 'x' of length 1" ):
14431451 xr .map_blocks (lambda x : x .isel (x = [1 ]), obj , template = obj ).compute ()
14441452 with pytest .raises (TypeError , match = r"must be a DataArray" ):
1445- xr .map_blocks (lambda x : x .isel (x = [1 ]), obj , template = (obj ,)).compute ()
1453+ xr .map_blocks (lambda x : x .isel (x = [1 ]), obj , template = (obj ,)).compute () # type: ignore[arg-type]
14461454 with pytest .raises (ValueError , match = r"map_blocks requires that one block" ):
14471455 xr .map_blocks (
14481456 lambda x : x .isel (x = [1 ]).assign_coords (x = 10 ), obj , template = obj .isel (x = [1 ])
@@ -1778,10 +1786,10 @@ def test_graph_manipulation():
17781786 for a , b in ((v , v2 ), (da , da2 ), (ds , ds2 )):
17791787 assert a .__dask_layers__ () != b .__dask_layers__ ()
17801788 assert len (a .__dask_layers__ ()) == len (b .__dask_layers__ ())
1781- assert a .__dask_graph__ ().keys () != b .__dask_graph__ ().keys ()
1782- assert len (a .__dask_graph__ ()) == len (b .__dask_graph__ ())
1783- assert a .__dask_graph__ ().layers .keys () != b .__dask_graph__ ().layers .keys ()
1784- assert len (a .__dask_graph__ ().layers ) == len (b .__dask_graph__ ().layers )
1789+ assert a .__dask_graph__ ().keys () != b .__dask_graph__ ().keys () # type: ignore[union-attr]
1790+ assert len (a .__dask_graph__ ()) == len (b .__dask_graph__ ()) # type: ignore[arg-type]
1791+ assert a .__dask_graph__ ().layers .keys () != b .__dask_graph__ ().layers .keys () # type: ignore[union-attr]
1792+ assert len (a .__dask_graph__ ().layers ) == len (b .__dask_graph__ ().layers ) # type: ignore[union-attr]
17851793
17861794 # Above we performed a slice operation; adding the two slices back together creates
17871795 # a diamond-shaped dependency graph, which in turn will trigger a collision in layer
0 commit comments