You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
61 lines
1.0 KiB
61 lines
1.0 KiB
from dask.context import globalmethod
|
|
import dask.array as da
|
|
import dask
|
|
|
|
|
|
def test_with_get():
|
|
var = [0]
|
|
|
|
def myget(dsk, keys, **kwargs):
|
|
var[0] = var[0] + 1
|
|
return dask.get(dsk, keys, **kwargs)
|
|
|
|
x = da.ones(10, chunks=(5,))
|
|
|
|
assert x.sum().compute() == 10
|
|
assert var[0] == 0
|
|
|
|
with dask.config.set(scheduler=myget):
|
|
assert x.sum().compute() == 10
|
|
assert var[0] == 1
|
|
|
|
# Make sure we've cleaned up
|
|
assert x.sum().compute() == 10
|
|
assert var[0] == 1
|
|
|
|
|
|
def foo():
|
|
return 'foo'
|
|
|
|
|
|
def bar():
|
|
return 'bar'
|
|
|
|
|
|
class Foo(object):
|
|
@globalmethod(key='f')
|
|
def f():
|
|
return 1
|
|
|
|
g = globalmethod(foo, key='g', falsey=bar)
|
|
|
|
|
|
def test_globalmethod():
|
|
x = Foo()
|
|
|
|
assert x.f() == 1
|
|
|
|
with dask.config.set(f=lambda: 2):
|
|
assert x.f() == 2
|
|
|
|
with dask.config.set(f=foo):
|
|
assert x.f is foo
|
|
assert x.f() == 'foo'
|
|
|
|
assert x.g is foo
|
|
assert x.g() == 'foo'
|
|
|
|
with dask.config.set(g=False):
|
|
assert x.g is bar
|
|
assert x.g() == 'bar'
|