ronantakizawa/codereview-bench huggingface dataset:
Repository: dask/dask
File: dask/tests/test_task_spec.py (Python)

Code before:
    try:
        return "-".join(args)
    except TypeError:
        return "-".join(map(str, args))


def func2(*args):
    return "=".join(args)


def func3(*args, **kwargs):
    return "+".join(args) + "//" + "+".join(f"{k}={v}" for k, v in kwargs.items())


def test_convert_legacy_dsk_skip_new():
    dsk = {
        "key-1": Task("key-1", func, "a", "b"),
    }
    converted = convert_legacy_graph(dsk)
    assert converted["key-1"] is dsk["key-1"]
    assert converted == dsk


def _assert_dsk_conversion(new_dsk):
    vals = [(k, v.key) for k, v in new_dsk.items()]
    assert all(v[0] == v[1] for v in vals), vals


def test_convert_legacy_dsk():
    def func(*args):
        return "-".join(args)

    def func2(*args):
        return "=".join(args)

    dsk = {
        "key-1": (func, "a", "b"),
        "key-2": (func2, "key-1", "c"),
        "key-3": (func, (func2, "c", "key-1"), "key-2"),
        "const": "foo",
        "key-4": [
            (func, "key-1", "b"),
            (func, "c", "key-2"),
            (func, "key-3", "f"),
            (func, "const", "bar"),
        ],
    }
    new_dsk = convert_legacy_graph(dsk)
    _assert_dsk_conversion(new_dsk)
    t1 = new_dsk["key-1"]
    assert isinstance(t1, Task)

Review comment (suggestion):
Thoughts on adding an equality sign here?

```suggestion
    assert repr(t) == "<Task 'key'=func('a', 'b')>"
```

Code after:
    try:
        return "-".join(args)
    except TypeError:
        return "-".join(map(str, args))


def func2(*args):
    return "=".join(args)


def func3(*args, **kwargs):
    return "+".join(args) + "//" + "+".join(f"{k}={v}" for k, v in kwargs.items())


def test_convert_legacy_dsk_skip_new():
    dsk = {
        "key-1": Task("key-1", func, "a", "b"),
    }
    converted = convert_legacy_graph(dsk)
    assert converted["key-1"] is dsk["key-1"]
    assert converted == dsk


def test_repr():
    t = Task("key", func, "a", "b")
    assert repr(t) == "<Task 'key' func('a', 'b')>"

    t = Task("nested", func2, t, t.ref())
    assert (
        repr(t) == "<Task 'nested' func2(<Task 'key' func('a', 'b')>, Alias(key->key))>"
    )

    def long_function_name_longer_even_longer(a, b):
        return a + b

    t = Task("long", long_function_name_longer_even_longer, t, t.ref())
    assert repr(t) == "<Task 'long' long_function_name_longer_even_longer(...)>"

    def use_kwargs(a, kwarg=None):
        return a + kwarg

    t = Task("kwarg", use_kwargs, "foo", kwarg="kwarg_value")
    assert repr(t) == "<Task 'kwarg' use_kwargs('foo', kwarg='kwarg_value')>"


def _assert_dsk_conversion(new_dsk):
    vals = [(k, v.key) for k, v in new_dsk.items()]
    assert all(v[0] == v[1] for v in vals), vals


def test_convert_legacy_dsk():
Repository: dask/dask
File: dask/array/tests/test_overlap.py (Python)

Code before:
        (2, (0, 1)),  # len(window shape) < len(axis)
        (2, None),  # len(window shape) < len(axis)
        (0, None),  # window_shape = 0
    ],
)
def test_sliding_window_errors(window_shape, axis):
    arr = da.zeros((4, 3))
    with pytest.raises(ValueError):
        sliding_window_view(arr, window_shape, axis)


def test_map_overlap_new_axis():
    arr = da.arange(6, chunks=2)
    assert arr.shape == (6,)
    assert arr.chunks == ((2, 2, 2),)

    actual = arr.map_overlap(lambda x: np.stack([x, x + 0.5]), depth=1, new_axis=[0])
    expected = np.stack([np.arange(6), np.arange(6) + 0.5])

    assert actual.chunks == ((1,), (2, 2, 2))
    # Shape and chunks aren't known until array is computed,
    # so don't expclitly check shape or chunks in assert_eq
    assert_eq(expected, actual, check_shape=False, check_chunks=False)


def test_overlap_now_blowing_up_graph():
    xr = pytest.importorskip("xarray")

    arr = xr.DataArray(
        da.random.random((2, 10, 5, 121), chunks=(1, 2, 1, 121)),
        dims=["x", "y", "z", "year"],
    )
    result = arr.rolling(dim={"year": 11}, center=True).construct("window_dim")
    dc = collections_to_dsk([result.data])
    assert len(dc) < 1000  # previously 3000


Review comment (suggestion):
```suggestion
def test_overlap_not_blowing_up_graph():
```

Code after:
        (2, (0, 1)),  # len(window shape) < len(axis)
        (2, None),  # len(window shape) < len(axis)
        (0, None),  # window_shape = 0
    ],
)
def test_sliding_window_errors(window_shape, axis):
    arr = da.zeros((4, 3))
    with pytest.raises(ValueError):
        sliding_window_view(arr, window_shape, axis)


def test_map_overlap_new_axis():
    arr = da.arange(6, chunks=2)
    assert arr.shape == (6,)
    assert arr.chunks == ((2, 2, 2),)

    actual = arr.map_overlap(lambda x: np.stack([x, x + 0.5]), depth=1, new_axis=[0])
    expected = np.stack([np.arange(6), np.arange(6) + 0.5])

    assert actual.chunks == ((1,), (2, 2, 2))
    # Shape and chunks aren't known until array is computed,
    # so don't expclitly check shape or chunks in assert_eq
    assert_eq(expected, actual, check_shape=False, check_chunks=False)


def test_overlap_not_blowing_up_graph():
    xr = pytest.importorskip("xarray")

    arr = xr.DataArray(
        da.random.random((2, 10, 5, 121), chunks=(1, 2, 1, 121)),
        dims=["x", "y", "z", "year"],
    )
    result = arr.rolling(dim={"year": 11}, center=True).construct("window_dim")
    dc = collections_to_dsk([result.data])
    assert len(dc) < 1000  # previously 3000
Repository: dask/dask
File: dask/array/_shuffle.py (Python)

Code before:
                )
            elif len(merge_keys) == 1:
                merges[(out_name,) + merge_suffix] = intermediates.pop(merge_keys[0])
            else:
                raise NotImplementedError

    output_chunks = []
    for i, c in enumerate(chunks):
        if i == axis:
            output_chunks.append(tuple(map(len, new_chunks)))
        else:
            output_chunks.append(c)

    layer = {**merges, **intermediates}
    return tuple(output_chunks), layer


def _getitem(obj, index):
    return getitem(obj, index[1])


def concatenate_arrays(arrs, sorter, axis):
    typ = type(arrs[0])
    take = take_lookup.dispatch(typ)
    concatenate = concatenate_lookup.dispatch(typ)
    return take(concatenate(arrs, axis=axis), np.argsort(sorter[1]), axis=axis)


def convert_key(key, chunk, axis):
    key = list(key)
    key.insert(axis, chunk)
    return tuple(key)


Review comment (suggestion):
```suggestion
    return take_lookup(
        concatenate_lookup(arrs, axis=axis),
        np.argsort(sorter[1]),
        axis=axis,
    )
```

Code after:
                )
            elif len(merge_keys) == 1:
                merges[(out_name,) + merge_suffix] = intermediates.pop(merge_keys[0])
            else:
                raise NotImplementedError

    output_chunks = []
    for i, c in enumerate(chunks):
        if i == axis:
            output_chunks.append(tuple(map(len, new_chunks)))
        else:
            output_chunks.append(c)

    layer = {**merges, **intermediates}
    return tuple(output_chunks), layer


def _getitem(obj, index):
    return getitem(obj, index[1])


def concatenate_arrays(arrs, sorter, axis):
    return take_lookup(
        concatenate_lookup.dispatch(type(arrs[0]))(arrs, axis=axis),
        np.argsort(sorter[1]),
        axis=axis,
    )


def convert_key(key, chunk, axis):
    key = list(key)
    key.insert(axis, chunk)
    return tuple(key)
Repository: dask/dask
File: dask/array/reductions.py (Python)

Code before:
        return quantiles[0].squeeze(axis=0)
    else:
        return np.concatenate(quantiles, axis=0)


@derived_from(np)
def nanquantile(
    a,
    q,
    axis=None,
    out=None,
    overwrite_input=False,
    method="linear",
    keepdims=False,
    *,
    weights=None,
    interpolation=None,
):
    """
    This works by automatically chunking the reduced axes to a single chunk
    and then calling ``numpy.nanquantile`` function across the remaining dimensions
    """
    if axis is None:
        if any(n_blocks > 1 for n_blocks in a.numblocks):
            raise NotImplementedError(
                "The da.quantile function only works along an axis.  "
                "The full algorithm is difficult to do in parallel"
            )
        else:
            axis = tuple(range(len(a.shape)))

    if not isinstance(axis, Iterable):
        axis = (axis,)

    axis = [ax + a.ndim if ax < 0 else ax for ax in axis]

    # rechunk if reduced axes are not contained in a single chunk
    if builtins.any(a.numblocks[ax] > 1 for ax in axis):
        a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})

    if NUMPY_GE_200:
        kwargs = {"weights": weights}
    else:
        kwargs = {}

    result = a.map_blocks(
        _custom_quantile,
        q=q,
        method=method,
        interpolation=interpolation,
        axis=axis,

Review comment (suggestion):

```suggestion
                "The da.nanquantile function only works along an axis.  "
```

Code after:
        return quantiles[0].squeeze(axis=0)
    else:
        return np.concatenate(quantiles, axis=0)


@derived_from(np)
def nanquantile(
    a,
    q,
    axis=None,
    out=None,
    overwrite_input=False,
    method="linear",
    keepdims=False,
    *,
    weights=None,
    interpolation=None,
):
    """
    This works by automatically chunking the reduced axes to a single chunk
    and then calling ``numpy.nanquantile`` function across the remaining dimensions
    """
    if axis is None:
        if any(n_blocks > 1 for n_blocks in a.numblocks):
            raise NotImplementedError(
                "The da.nanquantile function only works along an axis.  "
                "The full algorithm is difficult to do in parallel"
            )
        else:
            axis = tuple(range(len(a.shape)))

    if not isinstance(axis, Iterable):
        axis = (axis,)

    axis = [ax + a.ndim if ax < 0 else ax for ax in axis]

    # rechunk if reduced axes are not contained in a single chunk
    if builtins.any(a.numblocks[ax] > 1 for ax in axis):
        a = a.rechunk({ax: -1 if ax in axis else "auto" for ax in range(a.ndim)})

    if NUMPY_GE_200:
        kwargs = {"weights": weights}
    else:
        kwargs = {}

    result = a.map_blocks(
        _custom_quantile,
        q=q,
        method=method,
        interpolation=interpolation,
        axis=axis,
Repository: dask/dask
File: dask/_task_spec.py (Python)

Code before:
    def __repr__(self):
        return f"Alias({self.key}->{self.target})"

    def __eq__(self, value: object) -> bool:
        if not isinstance(value, Alias):
            return False
        if self.key != value.key:
            return False
        if self.key != value.key:
            return False
        return True


class DataNode(GraphNode):
    value: Any
    typ: type
    __slots__ = tuple(__annotations__)

    def __init__(self, key: Any, value: Any):
        if key is None:
            key = (type(value).__name__, next(_anom_count))
        self.key = key
        self.value = value
        self.typ = type(value)
        self._dependencies = _no_deps

    def copy(self):
        return DataNode(self.key, self.value)

    def __call__(self, values=()):
        return self.value

    def __repr__(self):
        return f"DataNode({self.key}, type={self.typ}, {self.value})"

    def __dask_tokenize__(self):
        from dask.base import tokenize

        return (type(self).__name__, tokenize(self.value))

    def __reduce__(self) -> str | tuple[Any, ...]:
        return DataNode, (self.key, self.value)

    def __eq__(self, value: object) -> bool:
        if not isinstance(value, DataNode):
            return False
        if self.value != value.value:
            return False
        return True

    def __sizeof__(self) -> int:

Review comment (suggestion):
I suspect this is private because it shouldn't be used when instantiating new `Task` objects? This might be worth an inline comment.

Code after:
            return _no_deps
        return set().union(*map(_get_dependencies, obj.values()))
    elif isinstance(obj, (list, tuple, frozenset, set)):
        if not obj:
            return _no_deps
        return set().union(*map(_get_dependencies, obj))
    return _no_deps


class Task(GraphNode):
    func: Callable
    args: tuple
    kwargs: dict
    _token: str | None
    _is_coro: bool | None
    _repr: str | None

    __slots__ = tuple(__annotations__)

    def __init__(
        self,
        key: Any,
        func: Callable,
        /,
        *args: Any,
        _dependencies: set | frozenset | None = None,
        **kwargs: Any,
    ):
        self.key = key
        self.func = func
        self.args = tuple(
            Alias(obj.key) if isinstance(obj, TaskRef) else obj for obj in args
        )
        self.kwargs = {
            k: Alias(v.key) if isinstance(v, TaskRef) else v for k, v in kwargs.items()
        }
        if _dependencies is None:
            _dependencies = set()
            for a in itertools.chain(self.args, self.kwargs.values()):
                if isinstance(a, GraphNode):
                    _dependencies.update(a.dependencies)
        if _dependencies:
            self._dependencies = frozenset(_dependencies)
        else:
            self._dependencies = _no_deps
        self._is_coro = None
        self._token = None
        self._repr = None

    def copy(self):
        return type(self)(
Repository: dask/dask
File: dask/bag/tests/test_bag.py (Python)

Code before:

    def count_reify(*args, **kwargs):
        nonlocal calls
        calls += 1
        return reify(*args, **kwargs)

    monkeypatch.setattr(dask.bag.core, "reify", count_reify)
    b = db.from_sequence([1, 2, 3, 4, 5, 6], npartitions=1)
    b2 = b.map(inc).map(inc).map(inc)

    [d] = b2.to_delayed()

    calls = 0
    d.compute()
    assert calls == 1

    calls = 0
    b = db.from_sequence([1, 2, 3, 4, 5, 6], npartitions=1)
    b2 = b.map(inc).map(inc).map(inc)
    [d] = b2.to_delayed()
    assert d.__dask_layers__() != b2.__dask_layers__()
    [d2] = b2.to_delayed(optimize_graph=False)
    assert dict(d2.dask) == dict(b2.dask)
    assert d2.__dask_layers__() == b2.__dask_layers__()
    assert d.compute() == d2.compute()
    assert calls == 1 + 3

    x = b2.sum()
    d = x.to_delayed()
    assert d.__dask_layers__() == x.__dask_layers__()
    d2 = x.to_delayed(optimize_graph=False)
    assert dict(d2.dask) == dict(x.dask)
    assert d2.__dask_layers__() == x.__dask_layers__()

    calls = 0
    assert d.compute() == d2.compute()
    assert calls == 3

    calls = 0
    [d] = b2.to_textfiles(str(tmpdir), compute=False)
    assert calls == 0


def test_from_delayed():
    from dask.delayed import delayed

    a, b, c = delayed([1, 2, 3]), delayed([4, 5, 6]), delayed([7, 8, 9])
    bb = from_delayed([a, b, c])
    assert bb.name == from_delayed([a, b, c]).name

    assert isinstance(bb, Bag)

Review comment (suggestion):
Alternative suggestion that drives the point home.

```suggestion
    calls = 0
    result = d.compute()
    assert calls == 1
    calls = 0
    result2 = d2.compute()
    assert calls == 3
    assert result == result2
```

Code after:

    def count_reify(*args, **kwargs):
        nonlocal calls
        calls += 1
        return reify(*args, **kwargs)

    monkeypatch.setattr(dask.bag.core, "reify", count_reify)
    b = db.from_sequence([1, 2, 3, 4, 5, 6], npartitions=1)
    b2 = b.map(inc).map(inc).map(inc)

    [d] = b2.to_delayed()

    calls = 0
    d.compute()
    assert calls == 1

    calls = 0
    b = db.from_sequence([1, 2, 3, 4, 5, 6], npartitions=1)
    b2 = b.map(inc).map(inc).map(inc)
    [d] = b2.to_delayed()
    assert d.__dask_layers__() != b2.__dask_layers__()
    [d2] = b2.to_delayed(optimize_graph=False)
    assert dict(d2.dask) == dict(b2.dask)
    assert d2.__dask_layers__() == b2.__dask_layers__()
    calls = 0
    result = d.compute()
    assert calls == 1
    calls = 0
    result2 = d2.compute()
    assert calls == 3
    assert result == result2

    x = b2.sum()
    d = x.to_delayed()
    assert d.__dask_layers__() == x.__dask_layers__()
    d2 = x.to_delayed(optimize_graph=False)
    assert dict(d2.dask) == dict(x.dask)
    assert d2.__dask_layers__() == x.__dask_layers__()

    calls = 0
    assert d.compute() == d2.compute()
    assert calls == 3

    calls = 0
    [d] = b2.to_textfiles(str(tmpdir), compute=False)
    assert calls == 0


def test_from_delayed():
    from dask.delayed import delayed
Repository: dask/dask
File: dask/dataframe/dask_expr/diagnostics/_explain.py (Python)

Code before:
        ]
    )

    graph.node(explain_info["name"], label)


def _add_graphviz_edges(explain_info, graph):
    name = explain_info["name"]
    for _, dep in explain_info["dependencies"]:
        graph.edge(dep, name)


def _explain_info(expr: Expr):
    return {
        "name": expr._name,
        "label": funcname(type(expr)),
        "details": _explain_details(expr),
        "dependencies": _explain_dependencies(expr),
    }


def _explain_details(expr: Expr):
    details = {"npartitions": expr.npartitions}

    if isinstance(expr, Merge):
        details["how"] = expr.how
    elif isinstance(expr, ReadParquet):
        details["path"] = expr.path
    elif isinstance(expr, Projection):
        columns = expr.operand("columns")
        details["ncolumns"] = len(columns) if isinstance(columns, list) else "Series"

    return details


def _explain_dependencies(expr: Expr) -> list[tuple[str, str]]:
    dependencies = []
    for i, operand in enumerate(expr.operands):
        if not isinstance(operand, Expr):
            continue
        param = expr._parameters[i] if i < len(expr._parameters) else ""
        dependencies.append((str(param), operand._name))
    return dependencies


Review comment (suggestion):
I believe explain should just work out of the box for arrays (and HLGs) but the nodes wouldn't be annotated. Still, it's something

Code after:
            "}>",
        ]
    )

    graph.node(explain_info["name"], label)


def _add_graphviz_edges(explain_info, graph):
    name = explain_info["name"]
    for _, dep in explain_info["dependencies"]:
        graph.edge(dep, name)


def _explain_info(expr: BaseExpr):
    return {
        "name": expr._name,
        "label": funcname(type(expr)),
        "details": _explain_details(expr),
        "dependencies": _explain_dependencies(expr),
    }


def _explain_details(expr: BaseExpr):

    if not isinstance(expr, Expr):
        return {}
    details = {"npartitions": expr.npartitions}

    if isinstance(expr, Merge):
        details["how"] = expr.how
    elif isinstance(expr, ReadParquet):
        details["path"] = expr.path
    elif isinstance(expr, Projection):
        columns = expr.operand("columns")
        details["ncolumns"] = len(columns) if isinstance(columns, list) else "Series"

    return details


def _explain_dependencies(expr: BaseExpr) -> list[tuple[str, str]]:
    dependencies = []
    for i, operand in enumerate(expr.operands):
        if not isinstance(operand, BaseExpr):
            continue
        param = expr._parameters[i] if i < len(expr._parameters) else ""
        dependencies.append((str(param), operand._name))
    return dependencies
Repository: dask/dask
File: dask/tests/test_expr.py (Python)

Code before:
    from distributed import Nanny
    from distributed.utils_test import gen_cluster

    @gen_cluster(client=True, Worker=Nanny, nthreads=[("", 1)])
    async def test(c, s, a):

        expr = MyExprCachedProperty(foo=1, bar=2)
        for _ in range(10):
            assert expr.baz == 3
            assert expr.cached_property == 3

        assert MyExprCachedProperty.called_cached_property is True

        rt = pickle.loads(pickle.dumps(expr))
        assert rt.cached_property == 3
        assert MyExprCachedProperty.called_cached_property is True

        # Expressions are singletons, i.e. this doesn't crash
        expr2 = MyExprCachedProperty(foo=1, bar=2)
        assert expr2.cached_property == 3

        # But this does
        expr3 = MyExprCachedProperty(foo=1, bar=3)
        with pytest.raises(RuntimeError):
            assert expr3.cached_property == 4

        def f(expr):
            # We want the cache to be part of the pickle, i.e. this is a
            # different process such that the type is reset and the property can
            # be accessed without side effects
            assert MyExprCachedProperty.called_cached_property is False
            assert expr.cached_property == 3
            assert MyExprCachedProperty.called_cached_property is False

        await c.submit(f, expr)

    test()


Review comment (performance):
The test block using pytest.raises() includes an assert for equality after the exception should be raised, which is redundant. Instead, simply accessing expr3.cached_property inside the block is sufficient.
```suggestion
            expr3.cached_property
```

Code after:
    from distributed import Nanny
    from distributed.utils_test import gen_cluster

    @gen_cluster(client=True, Worker=Nanny, nthreads=[("", 1)])
    async def test(c, s, a):

        expr = MyExprCachedProperty(foo=1, bar=2)
        for _ in range(10):
            assert expr.baz == 3
            assert expr.cached_property == 3

        assert MyExprCachedProperty.called_cached_property is True

        rt = pickle.loads(pickle.dumps(expr))
        assert rt.cached_property == 3
        assert MyExprCachedProperty.called_cached_property is True

        # Expressions are singletons, i.e. this doesn't crash
        expr2 = MyExprCachedProperty(foo=1, bar=2)
        assert expr2.cached_property == 3

        # But this does
        expr3 = MyExprCachedProperty(foo=1, bar=3)
        with pytest.raises(RuntimeError):
            expr3.cached_property

        def f(expr):
            # We want the cache to be part of the pickle, i.e. this is a
            # different process such that the type is reset and the property can
            # be accessed without side effects
            assert MyExprCachedProperty.called_cached_property is False
            assert expr.cached_property == 3
            assert MyExprCachedProperty.called_cached_property is False

        await c.submit(f, expr)

    test()