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

Code before:
def test_collect_env_none(monkeypatch):
    monkeypatch.setenv("DASK_FOO", "bar")
    config = collect([])
    assert config.get("foo") == "bar"


def test_get():
    d = {"x": 1, "y": {"a": 2}}

    assert get("x", config=d) == 1
    assert get("y.a", config=d) == 2
    assert get("y.b", 123, config=d) == 123
    with pytest.raises(KeyError):
        get("y.b", config=d)


def test_ensure_file(tmp_path):
    a = {"x": 1, "y": {"a": 1}}
    b = {"x": 123}

    source = tmp_path / "source.yaml"
    dest = tmp_path / "dest"
    destination = tmp_path / "dest" / "source.yaml"

    with source.open("w") as f:
        yaml.dump(a, f)

    ensure_file(source=source, destination=dest, comment=False)

    with destination.open() as f:
        result = yaml.safe_load(f)
    assert result == a

    # don't overwrite old config files
    with source.open("w") as f:
        yaml.dump(b, f)

    ensure_file(source=source, destination=dest, comment=False)

    with destination.open() as f:
        result = yaml.safe_load(f)
    assert result == a

    os.remove(destination)

    # Write again, now with comments
    ensure_file(source=source, destination=dest, comment=True)

    with destination.open() as f:
        text = f.read()
    assert "123" in text

Review comment (suggestion):
```suggestion
    source.write_text(yaml.dump(a))
```

Code after:
    a = {"x": 1, "y": {"a": 1}}
    b = {"x": 123}

    source = tmp_path / "source.yaml"
    dest = tmp_path / "dest"
    destination = tmp_path / "dest" / "source.yaml"
    source.write_text(yaml.dump(a))

    ensure_file(source=source, destination=dest, comment=False)

    result = yaml.safe_load(destination.read_text())
    assert result == a

    # don't overwrite old config files
    source.write_text(yaml.dump(b))

    ensure_file(source=source, destination=dest, comment=False)

    result = yaml.safe_load(destination.read_text())
    assert result == a

    os.remove(destination)

    # Write again, now with comments
    ensure_file(source=source, destination=dest, comment=True)

    text = destination.read_text()
    assert "123" in text

    result = yaml.safe_load(destination.read_text())
    assert not result


def test_set():
    with dask.config.set(abc=123):
        assert config["abc"] == 123
        with dask.config.set(abc=456):
            assert config["abc"] == 456
        assert config["abc"] == 123

    assert "abc" not in config

    with dask.config.set({"abc": 123}):
        assert config["abc"] == 123
    assert "abc" not in config

    with dask.config.set({"abc.x": 1, "abc.y": 2, "abc.z.a": 3}):
        assert config["abc"] == {"x": 1, "y": 2, "z": {"a": 3}}
    assert "abc" not in config

    d = {}
Repository: dask/dask
File: dask/tests/test_config.py (Python)

Code before:

def test_get():
    d = {"x": 1, "y": {"a": 2}}

    assert get("x", config=d) == 1
    assert get("y.a", config=d) == 2
    assert get("y.b", 123, config=d) == 123
    with pytest.raises(KeyError):
        get("y.b", config=d)


def test_ensure_file(tmp_path):
    a = {"x": 1, "y": {"a": 1}}
    b = {"x": 123}

    source = tmp_path / "source.yaml"
    dest = tmp_path / "dest"
    destination = tmp_path / "dest" / "source.yaml"

    with source.open("w") as f:
        yaml.dump(a, f)

    ensure_file(source=source, destination=dest, comment=False)

    with destination.open() as f:
        result = yaml.safe_load(f)
    assert result == a

    # don't overwrite old config files
    with source.open("w") as f:
        yaml.dump(b, f)

    ensure_file(source=source, destination=dest, comment=False)

    with destination.open() as f:
        result = yaml.safe_load(f)
    assert result == a

    os.remove(destination)

    # Write again, now with comments
    ensure_file(source=source, destination=dest, comment=True)

    with destination.open() as f:
        text = f.read()
    assert "123" in text

    with destination.open() as f:
        result = yaml.safe_load(f)
    assert not result


Review comment (suggestion):
```suggestion
    result = yaml.safe_load(destination.read_text())
```

Code after:
    destination = tmp_path / "dest" / "source.yaml"
    source.write_text(yaml.dump(a))

    ensure_file(source=source, destination=dest, comment=False)

    result = yaml.safe_load(destination.read_text())
    assert result == a

    # don't overwrite old config files
    source.write_text(yaml.dump(b))

    ensure_file(source=source, destination=dest, comment=False)

    result = yaml.safe_load(destination.read_text())
    assert result == a

    os.remove(destination)

    # Write again, now with comments
    ensure_file(source=source, destination=dest, comment=True)

    text = destination.read_text()
    assert "123" in text

    result = yaml.safe_load(destination.read_text())
    assert not result


def test_set():
    with dask.config.set(abc=123):
        assert config["abc"] == 123
        with dask.config.set(abc=456):
            assert config["abc"] == 456
        assert config["abc"] == 123

    assert "abc" not in config

    with dask.config.set({"abc": 123}):
        assert config["abc"] == 123
    assert "abc" not in config

    with dask.config.set({"abc.x": 1, "abc.y": 2, "abc.z.a": 3}):
        assert config["abc"] == {"x": 1, "y": 2, "z": {"a": 3}}
    assert "abc" not in config

    d = {}
    dask.config.set({"abc.x": 123}, config=d)
    assert d["abc"]["x"] == 123

    # Respect previous hypenation, if any, or new hypentation, if previous is not found
    d = {"e_f": 0, "g-h": 1}
Repository: dask/dask
File: continuous_integration/scripts/test_imports.sh (Shell)

Code before:


test_import () {
    echo "Create environment: python=$PYTHON_VERSION $1"
    # Create an empty environment
    mamba create -q -y -n test-imports -c conda-forge python=$PYTHON_VERSION packaging pyyaml fsspec toolz partd click cloudpickle importlib-metadata $1
    conda activate test-imports
    if [[ $1 =~ "distributed" ]]; then
        # dask[distributed] depends on the latest version of distributed
        python -m pip install git+https://github.com/dask/distributed
    fi
    python -m pip install -e .
    mamba list
    echo "python -c '$2'"
    python -c "$2"
    # Ensure that no non-deterministic objects are tokenized at init time,
    # which can prevent the library from being imported at all.
    echo "python -c '$2' (ensure deterministic)"
    DASK_TOKENIZE__ENSURE_DETERMINISTIC=True python -c "$2"
    conda deactivate
    mamba env remove -n test-imports
}

test_import ""                                "import dask, dask.base, dask.multiprocessing, dask.threaded, dask.optimization, dask.bag, dask.delayed, dask.graph_manipulation, dask.layers"
test_import "numpy"                           "import dask.array"
test_import "pandas"                          "import dask.dataframe"
test_import "bokeh"                           "import dask.diagnostics"
test_import "distributed"                     "import dask.distributed"


Review comment (suggestion):
I think this means `dask-expr` is now a requirement to `import dask.dataframe` (I could be missing something). Is that intentional? If so, we should update requirements too 

Code after:


test_import () {
    echo "Create environment: python=$PYTHON_VERSION $1"
    # Create an empty environment
    mamba create -q -y -n test-imports -c conda-forge python=$PYTHON_VERSION packaging pyyaml fsspec toolz partd click cloudpickle importlib-metadata $1
    conda activate test-imports
    if [[ $1 =~ "distributed" ]]; then
        # dask[distributed] depends on the latest version of distributed
        python -m pip install git+https://github.com/dask/distributed
    fi
    python -m pip install -e .
    mamba list
    echo "python -c '$2'"
    python -c "$2"
    # Ensure that no non-deterministic objects are tokenized at init time,
    # which can prevent the library from being imported at all.
    echo "python -c '$2' (ensure deterministic)"
    DASK_TOKENIZE__ENSURE_DETERMINISTIC=True python -c "$2"
    conda deactivate
    mamba env remove -n test-imports
}

test_import ""                                "import dask, dask.base, dask.multiprocessing, dask.threaded, dask.optimization, dask.bag, dask.delayed, dask.graph_manipulation, dask.layers"
test_import "numpy"                           "import dask.array"
test_import "pandas dask-expr"                "import dask.dataframe"
test_import "bokeh"                           "import dask.diagnostics"
test_import "distributed"                     "import dask.distributed"
Repository: dask/dask
File: dask/dataframe/__init__.py (Python)

Code before:
from __future__ import annotations

import importlib
import warnings

from packaging.version import Version


def _dask_expr_enabled() -> bool:
    import pandas as pd

    import dask

    use_dask_expr = dask.config.get("dataframe.query-planning")
    if use_dask_expr or use_dask_expr is None and Version(pd.__version__).major >= 2:
        try:
            import dask_expr  # noqa: F401
        except ImportError:
            raise ValueError("Must install dask-expr to activate query planning.")
        use_dask_expr = True
    return use_dask_expr


try:
    import dask.dataframe._pyarrow_compat
    from dask.base import compute
    from dask.dataframe import backends, dispatch, methods, rolling
    from dask.dataframe._testing import test_dataframe
    from dask.dataframe.core import (
        DataFrame,
        Index,
        Series,
        _Frame,
        map_partitions,
        repartition,
        to_datetime,
        to_timedelta,
    )
    from dask.dataframe.groupby import Aggregation
    from dask.dataframe.io import (
        demo,
        from_array,
        from_dask_array,
        from_delayed,
        from_dict,
        from_map,

Review comment (suggestion):
I _think_ this is what we want, i.e., try to use `dask-expr` if enabled **and** `pandas>=2` is installed:

```suggestion
        return True
    return False
```

Code after:
from __future__ import annotations

import importlib
import warnings

from packaging.version import Version


def _dask_expr_enabled() -> bool:
    import pandas as pd

    import dask

    use_dask_expr = dask.config.get("dataframe.query-planning")
    if (
        use_dask_expr is False
        or use_dask_expr is None
        and Version(pd.__version__).major < 2
    ):
        return False
    try:
        import dask_expr  # noqa: F401
    except ImportError:
        raise ValueError("Must install dask-expr to activate query planning.")
    return True


try:
    import dask.dataframe._pyarrow_compat
    from dask.base import compute
    from dask.dataframe import backends, dispatch, methods, rolling
    from dask.dataframe._testing import test_dataframe
    from dask.dataframe.core import (
        DataFrame,
        Index,
        Series,
        _Frame,
        map_partitions,
        repartition,
        to_datetime,
        to_timedelta,
    )
    from dask.dataframe.groupby import Aggregation
    from dask.dataframe.io import (
        demo,
        from_array,
Repository: dask/dask
File: docs/source/conf.py (Python)

Code before:
#
# dask documentation build configuration file, created by
# sphinx-quickstart on Sun Jan  4 08:58:22 2015.
#
# This file is execfile()d with the current directory set to its containing dir.
#
# Note that not all possible configuration values are present in this
# autogenerated file.
#
# All configuration values have a default; values that are commented out
# serve to show the default.

from __future__ import annotations

import os

# Build docs with option disabled until we move dask-expr over, othwise
# we run into a lot of circular issues that sphinx can't resolve properly
os.environ["DASK_DATAFRAME__QUERY_PLANNING"] = "false"

# Add any Sphinx extension module names here, as strings. They can be extensions
# coming with Sphinx (named 'sphinx.ext.*') or your custom ones.
import sys

import sphinx_autosummary_accessors

# -- General configuration -----------------------------------------------------

# If your documentation needs a minimal Sphinx version, state it here.
# needs_sphinx = '1.0'


# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
sys.path.insert(0, os.path.abspath("../../"))

extensions = [
    "sphinx.ext.autodoc",
    "sphinx.ext.doctest",
    "sphinx.ext.mathjax",
    "sphinx.ext.intersphinx",

Review comment (suggestion):
```suggestion
# Build docs with option disabled until we move dask-expr over, otherwise
```

Code after:
#
# dask documentation build configuration file, created by
# sphinx-quickstart on Sun Jan  4 08:58:22 2015.
#
# This file is execfile()d with the current directory set to its containing dir.
#
# Note that not all possible configuration values are present in this
# autogenerated file.
#
# All configuration values have a default; values that are commented out
# serve to show the default.

from __future__ import annotations

import os

# Build docs with option disabled until we move dask-expr over, otherwise
# we run into a lot of circular issues that sphinx can't resolve properly
os.environ["DASK_DATAFRAME__QUERY_PLANNING"] = "false"

# Add any Sphinx extension module names here, as strings. They can be extensions
# coming with Sphinx (named 'sphinx.ext.*') or your custom ones.
import sys

import sphinx_autosummary_accessors

# -- General configuration -----------------------------------------------------

# If your documentation needs a minimal Sphinx version, state it here.
# needs_sphinx = '1.0'


# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
sys.path.insert(0, os.path.abspath("../../"))

extensions = [
    "sphinx.ext.autodoc",
    "sphinx.ext.doctest",
    "sphinx.ext.mathjax",
    "sphinx.ext.intersphinx",
Repository: dask/dask
File: continuous_integration/gpuci/build.sh (Shell)

Code before:
rapids-logger "Activate conda env"
. /opt/conda/etc/profile.d/conda.sh
conda activate dask

rapids-logger "Install distributed"
python -m pip install git+https://github.com/dask/distributed

rapids-logger "Install dask"
python -m pip install --no-deps -e .

rapids-logger "Install pytest-timeout"
python -m pip install pytest-timeout

rapids-logger "Check Python version"
python --version

rapids-logger "Check conda environment"
conda info
conda config --show-sources
conda list --show-channel-urls

rapids-logger "Python py.test for dask (LEGACY)"
DASK_DATAFRAME__QUERY_PLANNING=False py.test $WORKSPACE -n 3 -v -m gpu --junitxml="$WORKSPACE/junit-dask-legacy.xml" --cov-config="$WORKSPACE/pyproject.toml" --cov=dask --cov-report=xml:"$WORKSPACE/dask-coverage-legacy.xml" --cov-report term

rapids-logger "Python py.test for dask (NEW)"
DASK_DATAFRAME__QUERY_PLANNING=True py.test $WORKSPACE -n 3 -v -m gpu --junitxml="$WORKSPACE/junit-dask.xml" --cov-config="$WORKSPACE/pyproject.toml" --cov=dask --cov-report=xml:"$WORKSPACE/dask-coverage.xml" --cov-report term


Review comment (suggestion):
@charlesbluca - How do you suggest we test both "legacy" and "dask-expr"? (We probably want to do this for 24.06, but can probably drop the legacy testing after that)

Code after:
rapids-logger "Activate conda env"
. /opt/conda/etc/profile.d/conda.sh
conda activate dask

rapids-logger "Install distributed"
python -m pip install git+https://github.com/dask/distributed

rapids-logger "Install dask"
python -m pip install --no-deps -e .

rapids-logger "Install pytest-timeout"
python -m pip install pytest-timeout

rapids-logger "Check Python version"
python --version

rapids-logger "Check conda environment"
conda info
conda config --show-sources
conda list --show-channel-urls

rapids-logger "Python py.test for dask-dataframe (LEGACY)"
DASK_DATAFRAME__QUERY_PLANNING=False py.test "$WORKSPACE/dask/dataframe" -n 3 -v -m gpu --junitxml="$WORKSPACE/junit-dask-legacy.xml" --cov-config="$WORKSPACE/pyproject.toml" --cov=dask --cov-report=xml:"$WORKSPACE/dask-coverage-legacy.xml" --cov-report term

rapids-logger "Python py.test for dask (NEW)"
DASK_DATAFRAME__QUERY_PLANNING=True py.test $WORKSPACE -n 3 -v -m gpu --junitxml="$WORKSPACE/junit-dask.xml" --cov-config="$WORKSPACE/pyproject.toml" --cov=dask --cov-report=xml:"$WORKSPACE/dask-coverage.xml" --cov-report term
Repository: dask/dask
File: dask/dataframe/tests/test_dataframe.py (Python)

Code before:
        divisions[0] -= 10
        ddf.divisions = tuple(divisions)
        assert_eq(pdf, ddf.enforce_runtime_divisions())

        # Setting an incorrect division boundary should
        # produce a `RuntimeError` in `compute`
        divisions[2] -= 10
        ddf.divisions = tuple(divisions)
        with pytest.raises(
            RuntimeError, match="`enforce_runtime_divisions` failed for partition 1"
        ):
            ddf.enforce_runtime_divisions().compute()


def test_preserve_ts_unit_in_meta_creation():
    pdf = pd.DataFrame(
        {
            "a": [1],
            "timestamp": pd.Series(
                [pd.Timestamp.utcnow()], dtype="datetime64[us, UTC]"
            ),
        }
    )
    df = dd.from_pandas(pdf, npartitions=1)
    dd.assert_eq(meta_nonempty(df._meta).dtypes, pdf.dtypes)
    dd.assert_eq(df, pdf)


def test_query_planning_config_warns():
    # Make sure dd._dask_expr_enabled() warns if the current
    # "dataframe.query-planning" config conflicts with the
    # global dd.DASK_EXPR_ENABLED setting.
    with dask.config.set({"dataframe.query-planning": not DASK_EXPR_ENABLED}):
        expect = "enabled" if dd.DASK_EXPR_ENABLED else "disabled"
        with pytest.warns(match=f"query planning is already {expect}"):
            dd._dask_expr_enabled()


Review comment (suggestion):
Minor non-blocking nit: `assert_eq` is already imported at the top and used throughout -- probably makes sense do the same here. 

```suggestion
    assert_eq(meta_nonempty(df._meta).dtypes, pdf.dtypes)
    assert_eq(df, pdf)
```

Code after:
        divisions[0] -= 10
        ddf.divisions = tuple(divisions)
        assert_eq(pdf, ddf.enforce_runtime_divisions())

        # Setting an incorrect division boundary should
        # produce a `RuntimeError` in `compute`
        divisions[2] -= 10
        ddf.divisions = tuple(divisions)
        with pytest.raises(
            RuntimeError, match="`enforce_runtime_divisions` failed for partition 1"
        ):
            ddf.enforce_runtime_divisions().compute()


def test_preserve_ts_unit_in_meta_creation():
    pdf = pd.DataFrame(
        {
            "a": [1],
            "timestamp": pd.Series(
                [pd.Timestamp.utcnow()], dtype="datetime64[us, UTC]"
            ),
        }
    )
    df = dd.from_pandas(pdf, npartitions=1)
    assert_eq(meta_nonempty(df._meta).dtypes, pdf.dtypes)
    assert_eq(df, pdf)


def test_query_planning_config_warns():
    # Make sure dd._dask_expr_enabled() warns if the current
    # "dataframe.query-planning" config conflicts with the
    # global dd.DASK_EXPR_ENABLED setting.
    with dask.config.set({"dataframe.query-planning": not DASK_EXPR_ENABLED}):
        expect = "enabled" if dd.DASK_EXPR_ENABLED else "disabled"
        with pytest.warns(match=f"query planning is already {expect}"):
            dd._dask_expr_enabled()
Repository: dask/dask
File: dask/array/tests/test_xarray.py (Python)

Code before:


def test_asarray_xarray_intersphinx_workaround():
    # test that the intersphinx workaround in https://github.com/pydata/xarray/issues/4279 works
    module = xr.DataArray.__module__
    try:
        xr.DataArray.__module__ = "xarray"
        y = da.asarray(xr.DataArray([1, 2, 3.0]))
        assert isinstance(y, da.Array)
        assert type(y._meta).__name__ == "ndarray"
        assert_eq(y, y)
    finally:
        xr.DataArray.__module__ = module


def test_fft():
    # Regression test for https://github.com/dask/dask/issues/9679
    coord = da.arange(8, chunks=-1)
    data = da.random.random((8, 8), chunks=-1) + 1
    x = xr.DataArray(data, coords={"x": coord, "y": coord}, dims=["x", "y"])
    result = da.fft.fft(x)
    expected = da.fft.fft(x.data)
    assert_eq(result, expected)


def test_polyfit_reshaping():
    arr = xr.DataArray(da.ones((10, 20, 30), chunks=(1, 5, 30)), dims=["z", "y", "x"])
    result = arr.polyfit("x", 1)
    assert result.polyfit_coefficients.chunks == ((2,), (1,) * 10, (5,) * 4)


Review comment (suggestion):
```suggestion
def test_polyfit_reshaping():
    # Regression test for https://github.com/pydata/xarray/issues/4554
```

Code after:


def test_asarray_xarray_intersphinx_workaround():
    # test that the intersphinx workaround in https://github.com/pydata/xarray/issues/4279 works
    module = xr.DataArray.__module__
    try:
        xr.DataArray.__module__ = "xarray"
        y = da.asarray(xr.DataArray([1, 2, 3.0]))
        assert isinstance(y, da.Array)
        assert type(y._meta).__name__ == "ndarray"
        assert_eq(y, y)
    finally:
        xr.DataArray.__module__ = module


def test_fft():
    # Regression test for https://github.com/dask/dask/issues/9679
    coord = da.arange(8, chunks=-1)
    data = da.random.random((8, 8), chunks=-1) + 1
    x = xr.DataArray(data, coords={"x": coord, "y": coord}, dims=["x", "y"])
    result = da.fft.fft(x)
    expected = da.fft.fft(x.data)
    assert_eq(result, expected)


def test_polyfit_reshaping():
    # Regression test for https://github.com/pydata/xarray/issues/4554
    arr = xr.DataArray(da.ones((10, 20, 30), chunks=(1, 5, 30)), dims=["z", "y", "x"])
    result = arr.polyfit("x", 1)
    assert result.polyfit_coefficients.chunks == ((2,), (1,) * 10, (5,) * 4)