<s> from itertools import chain <EOL> from dask . order import dfs , child_max , ndependents , order , inc , get_deps <EOL> def issorted ( L , reverse = False ) : <EOL> return sorted ( L , reverse = reverse ) == L <EOL> def f ( * args ) : <EOL> pass <EOL> def test_ordering_keeps_groups_together ( ) : <EOL> a , b , c = '<STR_LIT:abc>' <EOL> d = dict ( ( ( a , i ) , ( f , ) ) for i in range ( <NUM_LIT:4> ) ) <EOL> d . update ( { ( b , <NUM_LIT:0> ) : ( f , ( a , <NUM_LIT:0> ) , ( a , <NUM_LIT:1> ) ) , <EOL> ( b , <NUM_LIT:1> ) : ( f , ( a , <NUM_LIT:2> ) , ( a , <NUM_LIT:3> ) ) } ) <EOL> o = order ( d ) <EOL> assert abs ( o [ ( a , <NUM_LIT:0> ) ] - o [ ( a , <NUM_LIT:1> ) ] ) == <NUM_LIT:1> <EOL> assert abs ( o [ ( a , <NUM_LIT:2> ) ] - o [ ( a , <NUM_LIT:3> ) ] ) == <NUM_LIT:1> <EOL> d = dict ( ( ( a , i ) , ( f , ) ) for i in range ( <NUM_LIT:4> ) ) <EOL> d . update ( { ( b , <NUM_LIT:0> ) : ( f , ( a , <NUM_LIT:0> ) , ( a , <NUM_LIT:2> ) ) , <EOL> ( b , <NUM_LIT:1> ) : ( f , ( a , <NUM_LIT:1> ) , ( a , <NUM_LIT:3> ) ) } ) <EOL> o = order ( d ) <EOL> assert abs ( o [ ( a , <NUM_LIT:0> ) ] - o [ ( a , <NUM_LIT:2> ) ] ) == <NUM_LIT:1> <EOL> assert abs ( o [ ( a , <NUM_LIT:1> ) ] - o [ ( a , <NUM_LIT:3> ) ] ) == <NUM_LIT:1> <EOL> def test_prefer_broker_nodes ( ) : <EOL> """<STR_LIT>""" <EOL> a , b , c = '<STR_LIT:abc>' <EOL> dsk = { ( a , <NUM_LIT:0> ) : ( f , ) , ( a , <NUM_LIT:1> ) : ( f , ) , <EOL> ( b , <NUM_LIT:0> ) : ( f , ( a , <NUM_LIT:0> ) ) , ( b , <NUM_LIT:1> ) : ( f , ( a , <NUM_LIT:1> ) ) , ( b , <NUM_LIT:2> ) : ( f , ( a , <NUM_LIT:1> ) ) } <EOL> dependencies , dependents = get_deps ( dsk ) <EOL> nd = ndependents ( dependencies , dependents ) <EOL> cm = child_max ( dependencies , dependents , nd ) <EOL> o = order ( dsk ) <EOL> assert o [ ( a , <NUM_LIT:1> ) ] < o [ ( a , <NUM_LIT:0> ) ] <EOL> dsk = { ( a , <NUM_LIT:0> ) : ( f , ) , ( a , <NUM_LIT:1> ) : ( f , ) , <EOL> ( b , <NUM_LIT:0> ) : ( f , ( a , <NUM_LIT:0> ) ) , ( b , <NUM_LIT:1> ) : ( f , ( a , <NUM_LIT:1> ) ) , ( b , <NUM_LIT:2> ) : ( f , ( a , <NUM_LIT:0> ) ) } <EOL> o = order ( dsk ) <EOL> assert o [ ( a , <NUM_LIT:1> ) ] > o [ ( a , <NUM_LIT:0> ) ] <EOL> def test_base_of_reduce_preferred ( ) : <EOL> """<STR_LIT>""" <EOL> dsk = dict ( ( ( '<STR_LIT:a>' , i ) , ( f , ( '<STR_LIT:a>' , i - <NUM_LIT:1> ) , ( '<STR_LIT:b>' , i ) ) ) for i in [ <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] ) <EOL> dsk [ ( '<STR_LIT:a>' , <NUM_LIT:0> ) ] = ( f , ( '<STR_LIT:b>' , <NUM_LIT:0> ) ) <EOL> dsk . update ( dict ( ( ( '<STR_LIT:b>' , i ) , ( f , '<STR_LIT:c>' , <NUM_LIT:1> ) ) for i in [ <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] ) ) <EOL> dsk [ '<STR_LIT:c>' ] = <NUM_LIT:1> <EOL> o = order ( dsk ) <EOL> assert o == { ( '<STR_LIT:a>' , <NUM_LIT:3> ) : <NUM_LIT:0> , <EOL> ( '<STR_LIT:a>' , <NUM_LIT:2> ) : <NUM_LIT:1> , <EOL> ( '<STR_LIT:a>' , <NUM_LIT:1> ) : <NUM_LIT:2> , <EOL> ( '<STR_LIT:a>' , <NUM_LIT:0> ) : <NUM_LIT:3> , <EOL> ( '<STR_LIT:b>' , <NUM_LIT:0> ) : <NUM_LIT:4> , <EOL> '<STR_LIT:c>' : <NUM_LIT:5> , <EOL> ( '<STR_LIT:b>' , <NUM_LIT:1> ) : <NUM_LIT:6> , <EOL> ( '<STR_LIT:b>' , <NUM_LIT:2> ) : <NUM_LIT:7> , <EOL> ( '<STR_LIT:b>' , <NUM_LIT:3> ) : <NUM_LIT:8> } <EOL> assert min ( [ ( '<STR_LIT:b>' , i ) for i in [ <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] ] , key = o . get ) == ( '<STR_LIT:b>' , <NUM_LIT:0> ) <EOL> def test_deep_bases_win_over_dependents ( ) : <EOL> """<STR_LIT>""" <EOL> dsk = { '<STR_LIT:a>' : ( f , '<STR_LIT:b>' , '<STR_LIT:c>' , '<STR_LIT:d>' ) , '<STR_LIT:b>' : ( f , '<STR_LIT:d>' , '<STR_LIT:e>' ) , '<STR_LIT:c>' : ( f , '<STR_LIT:d>' ) , '<STR_LIT:d>' : <NUM_LIT:1> , <EOL> '<STR_LIT:e>' : <NUM_LIT:2> } <EOL> dependencies , dependents = get_deps ( dsk ) <EOL> nd = ndependents ( dependencies , dependents ) <EOL> cm = child_max ( dependencies , dependents , nd ) <EOL> o = order ( dsk ) <EOL> assert o [ '<STR_LIT:d>' ] < o [ '<STR_LIT:e>' ] <EOL> assert o [ '<STR_LIT:d>' ] < o [ '<STR_LIT:b>' ] or o [ '<STR_LIT:d>' ] < o [ '<STR_LIT:c>' ] <EOL> def test_prefer_deep ( ) : <EOL> """<STR_LIT>""" <EOL> dsk = { '<STR_LIT:a>' : <NUM_LIT:1> , '<STR_LIT:b>' : ( f , '<STR_LIT:a>' ) , '<STR_LIT:c>' : ( f , '<STR_LIT:b>' ) , <EOL> '<STR_LIT:x>' : <NUM_LIT:1> , '<STR_LIT:y>' : ( f , '<STR_LIT:x>' ) } <EOL> o = order ( dsk ) <EOL> assert o == { '<STR_LIT:c>' : <NUM_LIT:0> , '<STR_LIT:b>' : <NUM_LIT:1> , '<STR_LIT:a>' : <NUM_LIT:2> , '<STR_LIT:y>' : <NUM_LIT:3> , '<STR_LIT:x>' : <NUM_LIT:4> } <EOL> def test_stacklimit ( ) : <EOL> dsk = dict ( ( '<STR_LIT>' % ( i + <NUM_LIT:1> ) , ( inc , '<STR_LIT>' % i ) ) for i in range ( <NUM_LIT> ) ) <EOL> dependencies , dependents = get_deps ( dsk ) <EOL> scores = dict . fromkeys ( dsk , <NUM_LIT:1> ) <EOL> child_max ( dependencies , dependents , scores ) <EOL> ndependents ( dependencies , dependents ) <EOL> def test_ndependents ( ) : <EOL> from operator import add <EOL> a , b , c = '<STR_LIT:abc>' <EOL> dsk = dict ( chain ( ( ( ( a , i ) , i * <NUM_LIT:2> ) for i in range ( <NUM_LIT:5> ) ) , <EOL> ( ( ( b , i ) , ( add , i , ( a , i ) ) ) for i in range ( <NUM_LIT:5> ) ) , <EOL> ( ( ( c , i ) , ( add , i , ( b , i ) ) ) for i in range ( <NUM_LIT:5> ) ) ) ) <EOL> result = ndependents ( * get_deps ( dsk ) ) <EOL> expected = dict ( chain ( ( ( ( a , i ) , <NUM_LIT:3> ) for i in range ( <NUM_LIT:5> ) ) , <EOL> ( ( ( b , i ) , <NUM_LIT:2> ) for i in range ( <NUM_LIT:5> ) ) , <EOL> ( ( ( c , i ) , <NUM_LIT:1> ) for i in range ( <NUM_LIT:5> ) ) ) ) <EOL> assert result == expected <EOL> dsk = { a : <NUM_LIT:1> , b : <NUM_LIT:1> } <EOL> deps = get_deps ( dsk ) <EOL> assert ndependents ( * deps ) == dsk <EOL> dsk = { a : <NUM_LIT:1> , b : ( add , a , <NUM_LIT:1> ) , c : ( add , b , a ) } <EOL> assert ndependents ( * get_deps ( dsk ) ) == { a : <NUM_LIT:4> , b : <NUM_LIT:2> , c : <NUM_LIT:1> } <EOL> dsk = { a : <NUM_LIT:1> , b : a , c : b } <EOL> deps = get_deps ( dsk ) <EOL> assert ndependents ( * deps ) == { a : <NUM_LIT:3> , b : <NUM_LIT:2> , c : <NUM_LIT:1> } </s>
