Repository: dask/dask
File: dask/tests/test_optimization.py (Python)

Code before:

    d = {
        "w": (inc, "x"),
        "x": (inc, "y"),
        "y": (inc, "z"),
        "z": (add, "a", "b"),
        "a": 1,
        "b": 2,
    }
    keys = ["x", "z"]
    assert fuse(d, keys, rename_keys=False) == with_deps(
        {"w": (inc, "x"), "x": (inc, (inc, "z")), "z": (add, "a", "b"), "a": 1, "b": 2}
    )
    assert fuse(d, keys, rename_keys=True) == with_deps(
        {
            "w": (inc, "y-x"),
            "y-x": (inc, (inc, "z")),
            "z": (add, "a", "b"),
            "a": 1,
            "b": 2,
            "x": "y-x",
        }
    )


def test_inline():
    d = {"a": 1, "b": (inc, "a"), "c": (inc, "b"), "d": (add, "a", "c")}
    assert inline(d) == {"a": 1, "b": (inc, 1), "c": (inc, "b"), "d": (add, 1, "c")}
    assert inline(d, ["a", "b", "c"]) == {
        "a": 1,
        "b": (inc, 1),
        "c": (inc, (inc, 1)),
        "d": (add, 1, (inc, (inc, 1))),
    }
    d = {"x": 1, "y": (inc, "x"), "z": (add, "x", "y")}
    assert inline(d) == {"x": 1, "y": (inc, 1), "z": (add, 1, "y")}
    assert inline(d, keys="y") == {"x": 1, "y": (inc, 1), "z": (add, 1, (inc, 1))}
    assert inline(d, keys="y", inline_constants=False) == {
        "x": 1,
        "y": (inc, "x"),
        "z": (add, "x", (inc, "x")),
    }

    d = {"a": 1, "b": "a", "c": "b", "d": ["a", "b", "c"], "e": (add, (len, "d"), "a")}
    assert inline(d, "d") == {
        "a": 1,
        "b": 1,
        "c": 1,
        "d": [1, 1, 1],
        "e": (add, (len, [1, 1, 1]), 1),
    }

Review comment (bug):
```suggestion
def test_donot_substitute_same_key_multiple_times() -> None:
```
If you add `-> None` to new tests mypy will be able to check them as well. I think it's good practice to just do it if you are worrying about the type hints. 
It helps catch the edge cases that aren't always obvious when just type hinting the functions.

Code after:

    d = {
        "w": (inc, "x"),
        "x": (inc, "y"),
        "y": (inc, "z"),
        "z": (add, "a", "b"),
        "a": 1,
        "b": 2,
    }
    keys = ["x", "z"]
    assert fuse(d, keys, rename_keys=False) == with_deps(
        {"w": (inc, "x"), "x": (inc, (inc, "z")), "z": (add, "a", "b"), "a": 1, "b": 2}
    )
    assert fuse(d, keys, rename_keys=True) == with_deps(
        {
            "w": (inc, "y-x"),
            "y-x": (inc, (inc, "z")),
            "z": (add, "a", "b"),
            "a": 1,
            "b": 2,
            "x": "y-x",
        }
    )


def test_donot_substitute_same_key_multiple_times():
    already_called = False

    def inc_only_once(x):
        nonlocal already_called
        if already_called:
            raise RuntimeError
        already_called = True
        return x + 1

    # This is the graph topology of a Z.T @ Z after blockwise fusion as given in
    # https://github.com/dask/dask/issues/10645
    dsk = {
        # Note that there is some logic in there that actually checks the type
        # of this and an integer is inlined while an array is not. However, data
        # keys are also flagged as forbidden to be fused so this should not be a
        # problem
        "A": 42,  # Array
        "B": (inc_only_once, "A"),  # Add
        "C": (add, "B", "B"),  # matmul
        "D": (inc, "C"),  # getitem
    }

    # fuse accepts dependencies to avoid having to recompute it. However,
    # internally it actually requires it to be in the format of dict[Key,
    # list[Key]] which is rarely used elsewhere since most applications are