1- from typing import Optional , Tuple , Union
1+ from typing import List , Optional , Tuple , Union
22
33import numpy as np
44import pyarrow as pa
@@ -24,7 +24,7 @@ def test_soma_sparse_nd_array_ok_no_storage():
2424)
2525@pytest .mark .parametrize ("element_type" , NDARRAY_ARROW_TYPES_SUPPORTED )
2626def test_soma_sparse_nd_array_create_ok (
27- tmp_path , shape : Tuple [int , ...], element_type : pa .DataType
27+ tmp_path , shape : Union [ Tuple [int , ...], List [ int ] ], element_type : pa .DataType
2828):
2929 """
3030 Test all cases we expect "create" to succeed.
@@ -250,7 +250,7 @@ def test_soma_sparse_nd_array_read_write_sparse_tensor(
250250 assert a .shape == shape
251251
252252 # make a random sample in the desired format
253- data = create_random_tensor (format , shape , np .float64 )
253+ data = create_random_tensor (format , shape , np .dtype ( np . float64 ) )
254254 a .write_sparse_tensor (data )
255255 del a
256256
@@ -271,7 +271,7 @@ def test_soma_sparse_nd_array_read_write_table(
271271 assert a .shape == shape
272272
273273 # make a random sample in the desired format
274- data = create_random_tensor ("table" , shape , np .float32 )
274+ data = create_random_tensor ("table" , shape , np .dtype ( np . float32 ) )
275275 a .write_table (data )
276276 del a
277277
@@ -282,7 +282,15 @@ def test_soma_sparse_nd_array_read_write_table(
282282 assert tables_are_same_value (data , t )
283283
284284
285- @pytest .mark .parametrize ("dtype" , [np .float32 , np .float64 , np .int32 , np .int64 ])
285+ @pytest .mark .parametrize (
286+ "dtype" ,
287+ [
288+ np .dtype (np .float32 ),
289+ np .dtype (np .float64 ),
290+ np .dtype (np .int32 ),
291+ np .dtype (np .int64 ),
292+ ],
293+ )
286294@pytest .mark .parametrize ("shape" , [(1 ,), (23 , 14 ), (35 , 3 , 2 ), (8 , 4 , 2 , 30 )])
287295def test_soma_sparse_nd_array_read_as_pandas (
288296 tmp_path , dtype : np .dtype , shape : Tuple [int , ...]
@@ -378,7 +386,7 @@ def test_soma_sparse_nd_array_nnz(tmp_path):
378386 assert a .nnz == 0
379387
380388 t : pa .SparseCOOTensor = create_random_tensor (
381- "coo" , a .shape , pa .int32 ().to_pandas_dtype (), 0.1
389+ "coo" , a .shape , np . dtype ( pa .int32 ().to_pandas_dtype () ), 0.1
382390 )
383391 a .write_sparse_tensor (t )
384392 with pytest .raises (NotImplementedError ):
0 commit comments