5
5
import sys
6
6
from contextlib import suppress
7
7
from textwrap import dedent
8
+ from typing import Any
8
9
9
10
import numpy as np
10
11
import pandas as pd
@@ -94,13 +95,14 @@ def test_copy(self):
94
95
self .assertLazyAndIdentical (self .eager_var , self .lazy_var .copy (deep = True ))
95
96
96
97
def test_chunk (self ):
97
- for chunks , expected in [
98
+ test_cases : list [ tuple [ int | dict [ str , Any ], tuple [ tuple [ int , ...], ...]]] = [
98
99
({}, ((2 , 2 ), (2 , 2 , 2 ))),
99
100
(3 , ((3 , 1 ), (3 , 3 ))),
100
101
({"x" : 3 , "y" : 3 }, ((3 , 1 ), (3 , 3 ))),
101
102
({"x" : 3 }, ((3 , 1 ), (2 , 2 , 2 ))),
102
103
({"x" : (3 , 1 )}, ((3 , 1 ), (2 , 2 , 2 ))),
103
- ]:
104
+ ]
105
+ for chunks , expected in test_cases :
104
106
rechunked = self .lazy_var .chunk (chunks )
105
107
assert rechunked .chunks == expected
106
108
self .assertLazyAndIdentical (self .eager_var , rechunked )
@@ -258,7 +260,7 @@ def test_missing_methods(self):
258
260
with pytest .raises (NotImplementedError , match = "dask" ):
259
261
v .argsort ()
260
262
with pytest .raises (NotImplementedError , match = "dask" ):
261
- v [0 ].item ()
263
+ v [0 ].item () # type: ignore[attr-defined]
262
264
263
265
def test_univariate_ufunc (self ):
264
266
u = self .eager_var
@@ -298,7 +300,7 @@ def test_persist(self):
298
300
299
301
(v2 ,) = dask .persist (v )
300
302
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]
302
304
assert v2 .__dask_keys__ () == v .__dask_keys__ ()
303
305
assert dask .is_dask_collection (v )
304
306
assert dask .is_dask_collection (v2 )
@@ -345,7 +347,9 @@ def setUp(self):
345
347
)
346
348
347
349
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
+ ] = [
349
353
({}, ((2 , 2 ), (2 , 2 , 2 ))),
350
354
(3 , ((3 , 1 ), (3 , 3 ))),
351
355
({"x" : 3 , "y" : 3 }, ((3 , 1 ), (3 , 3 ))),
@@ -354,7 +358,8 @@ def test_chunk(self) -> None:
354
358
({"x" : "16B" }, ((1 , 1 , 1 , 1 ), (2 , 2 , 2 ))),
355
359
("16B" , ((1 , 1 , 1 , 1 ), (1 ,) * 6 )),
356
360
("16MB" , ((4 ,), (6 ,))),
357
- ]:
361
+ ]
362
+ for chunks , expected in test_cases :
358
363
# Test DataArray
359
364
rechunked = self .lazy_array .chunk (chunks )
360
365
assert rechunked .chunks == expected
@@ -367,7 +372,7 @@ def test_chunk(self) -> None:
367
372
lazy_dataset = self .lazy_array .to_dataset ()
368
373
eager_dataset = self .eager_array .to_dataset ()
369
374
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]
371
376
372
377
# Dataset.chunks has a different return type to DataArray.chunks - see issue #5843
373
378
assert rechunked .chunks == expected_chunksizes
@@ -601,11 +606,12 @@ def test_reindex(self):
601
606
u = self .eager_array .assign_coords (y = range (6 ))
602
607
v = self .lazy_array .assign_coords (y = range (6 ))
603
608
604
- for kwargs in [
609
+ kwargs_list : list [ dict [ str , Any ]] = [
605
610
{"x" : [2 , 3 , 4 ]},
606
611
{"x" : [1 , 100 , 2 , 101 , 3 ]},
607
612
{"x" : [2.5 , 3 , 3.5 ], "y" : [2 , 2.5 , 3 ]},
608
- ]:
613
+ ]
614
+ for kwargs in kwargs_list :
609
615
expected = u .reindex (** kwargs )
610
616
actual = v .reindex (** kwargs )
611
617
self .assertLazyAndAllClose (expected , actual )
@@ -666,7 +672,9 @@ def test_stack(self):
666
672
data = da .random .normal (size = (2 , 3 , 4 ), chunks = (1 , 3 , 4 ))
667
673
arr = DataArray (data , dims = ("w" , "x" , "y" ))
668
674
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
+ )
670
678
expected = DataArray (data .reshape (2 , - 1 ), {"z" : z }, dims = ["w" , "z" ])
671
679
assert stacked .data .chunks == expected .data .chunks
672
680
self .assertLazyAndEqual (expected , stacked )
@@ -1167,10 +1175,10 @@ def returns_numpy(darray):
1167
1175
xr .map_blocks (returns_numpy , map_da )
1168
1176
1169
1177
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]
1171
1179
1172
1180
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]
1174
1182
1175
1183
def really_bad_func (darray ):
1176
1184
raise ValueError ("couldn't do anything." )
@@ -1442,7 +1450,7 @@ def test_map_blocks_errors_bad_template(obj):
1442
1450
with pytest .raises (ValueError , match = r"Received dimension 'x' of length 1" ):
1443
1451
xr .map_blocks (lambda x : x .isel (x = [1 ]), obj , template = obj ).compute ()
1444
1452
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]
1446
1454
with pytest .raises (ValueError , match = r"map_blocks requires that one block" ):
1447
1455
xr .map_blocks (
1448
1456
lambda x : x .isel (x = [1 ]).assign_coords (x = 10 ), obj , template = obj .isel (x = [1 ])
@@ -1778,10 +1786,10 @@ def test_graph_manipulation():
1778
1786
for a , b in ((v , v2 ), (da , da2 ), (ds , ds2 )):
1779
1787
assert a .__dask_layers__ () != b .__dask_layers__ ()
1780
1788
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]
1785
1793
1786
1794
# Above we performed a slice operation; adding the two slices back together creates
1787
1795
# a diamond-shaped dependency graph, which in turn will trigger a collision in layer
0 commit comments