Repository: dask/dask
File: dask/dataframe/dask_expr/_expr.py (Python)

Code before:
            t = _expr._task(subname, subname[1])

            assert t.key == subname
            internal_tasks.append(t)
        return Task.fuse(*internal_tasks, key=name)  # type: ignore

    @staticmethod
    def _execute_internal_graph(internal_tasks, dependencies, outkey):
        cache = dict(dependencies)
        res = execute_graph(internal_tasks, cache=cache, keys=[outkey])
        return res[outkey]


# Used for sorting with None
@functools.total_ordering
class MinType:
    def __le__(self, other):
        return True


def determine_column_projection(
    expr: Expr,
    parent: Expr,
    dependents: dict[str, Collection[weakref.ref[BaseExpr]]],
    additional_columns: list | None = None,
):
    if isinstance(parent, Index):
        column_union = []
    else:
        column_union = parent.columns.copy()
    parents: list[Expr]
    parents = [
        x() for x in dependents[expr._name] if isinstance(x(), Expr)  # type: ignore
    ]

    seen = set()
    for p in parents:
        if p._name in seen:
            continue
        seen.add(p._name)
        column_union.extend(p._projection_columns)

    if additional_columns is not None:
        column_union.extend(flatten(additional_columns, container=list))

    # We can end up with MultiIndex columns from groupby ops, needs to be
    # accounted for in the sort
    flattened_columns = set(column_union)
    try:
        column_union = sorted(flattened_columns)
    except TypeError:

Review comment (suggestion):
I can't remember mypy's rules offhand (or whether we're using it), but do we need a return type for it to check the body of the function? Maybe not, since you have a `type: ignore` down below, in which case ignore this.

Code after:
            t = _expr._task(subname, subname[1])

            assert t.key == subname
            internal_tasks.append(t)
        return Task.fuse(*internal_tasks, key=name)  # type: ignore

    @staticmethod
    def _execute_internal_graph(internal_tasks, dependencies, outkey):
        cache = dict(dependencies)
        res = execute_graph(internal_tasks, cache=cache, keys=[outkey])
        return res[outkey]


# Used for sorting with None
@functools.total_ordering
class MinType:
    def __le__(self, other):
        return True


def determine_column_projection(
    expr: Expr,
    parent: Expr,
    dependents: dict[str, Collection[weakref.ref[BaseExpr]]],
    additional_columns: list | None = None,
) -> object:
    if isinstance(parent, Index):
        column_union = []
    else:
        column_union = parent.columns.copy()
    parents: list[Expr]
    parents = [inst for x in dependents[expr._name] if isinstance((inst := x()), Expr)]

    seen = set()
    for p in parents:
        if p._name in seen:
            continue
        seen.add(p._name)
        column_union.extend(p._projection_columns)

    if additional_columns is not None:
        column_union.extend(flatten(additional_columns, container=list))

    # We can end up with MultiIndex columns from groupby ops, needs to be
    # accounted for in the sort
    flattened_columns = set(column_union)
    try:
        column_union = sorted(flattened_columns)
    except TypeError:
        # mixed type columns
        column_union = _sort_mixed(pd.Index(list(flattened_columns))).tolist()