Repository: pypa/pip
File: src/pip/_internal/utils/filesystem.py (Python)

Code before:
def directory_size(path: str) -> int | float:
    size = 0.0
    for root, _dirs, files in os.walk(path):
        for filename in files:
            file_path = os.path.join(root, filename)
            size += file_size(file_path)
    return size


def format_directory_size(path: str) -> str:
    return format_size(directory_size(path))


def copy_directory_permissions(directory: str, target_file: BinaryIO) -> None:
    mode = (
        os.stat(directory).st_mode & 0o666  # select read/write permissions of directory
        | 0o600  # set owner read/write permissions
    )
    # Change permissions only if there is no risk of following a symlink.
    if os.chmod in os.supports_fd:
        os.chmod(target_file.fileno(), mode)
    elif os.chmod in os.supports_follow_symlinks:
        os.chmod(target_file.name, mode, follow_symlinks=False)


def subdirs_without_files(path: str) -> Generator[Path]:
    """Yields every subdirectory of +path+ that has no files under it."""

    directories = []
    non_empty = set()

    for root_str, _, filenames in os.walk(Path(path).resolve()):
        root = Path(root_str)
        if filenames:
            # This directory contains a file, so mark it and all of its
            # parent directories as non-empty.
            # The last item in root.parents is ".", so we ignore it.
            #
            # Wrapping this in `list()` is only needed for Python 3.9.
            non_empty.update(list(root.parents)[:-1])
            non_empty.add(root)
        directories.append(root)

    for d in sorted(directories, reverse=True):
        if d not in non_empty:
            yield d


def subdirs_without_wheels(path: str) -> Generator[Path]:
    """Yields every subdirectory of +path+ that has no .whl files under it."""


Review comment (security):
There is a fair bit of duplication. IMO, it'd nicer to extract the common logic into a generic helper function:

```diff
diff --git a/src/pip/_internal/utils/filesystem.py b/src/pip/_internal/utils/filesystem.py
index 014d1a271..52eb78f42 100644
--- a/src/pip/_internal/utils/filesystem.py
+++ b/src/pip/_internal/utils/filesystem.py
@@ -9,7 +9,7 @@ from collections.abc import Generator
 from contextlib import contextmanager
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import Any, BinaryIO, cast
+from typing import Any, BinaryIO, Callable, cast
 
 from pip._internal.utils.compat import get_path_uid
 from pip._internal.utils.misc import format_size
@@ -165,47 +165,39 @@ def copy_directory_permissions(directory: str, target_file: BinaryIO) -> None:
         os.chmod(target_file.name, mode, follow_symlinks=False)
 
 
-def subdirs_without_files(path: str) -> Generator[Path]:
-    """Yields every subdirectory of +path+ that has no files under it."""
+def _subdirs_without_generic(
+    path: str, predicate: Callable[[str, list[str]], bool]
+) -> Generator[Path]:
+    """Yields every subdirectory of +path+ that has no files matching the
+    predicate under it."""
 
     directories = []
-    non_empty = set()
+    excluded = set()
 
     for root_str, _, filenames in os.walk(Path(path).resolve()):
         root = Path(root_str)
-        if filenames:
-            # This directory contains a file, so mark it and all of its
-            # parent directories as non-empty.
+        if predicate(root_str, filenames):
+            # This directory should be excluded, so exclude it and all of its
+            # parent directories.
             # The last item in root.parents is ".", so we ignore it.
             #
             # Wrapping this in `list()` is only needed for Python 3.9.
-            non_empty.update(list(root.parents)[:-1])
-            non_empty.add(root)
+            excluded.update(list(root.parents)[:-1])
+            excluded.add(root)
         directories.append(root)
 
     for d in sorted(directories, reverse=True):
-        if d not in non_empty:
+        if d not in excluded:
             yield d
 
 
-def subdirs_without_wheels(path: str) -> Generator[Path]:
-    """Yields every subdirectory of +path+ that has no .whl files under it."""
+def subdirs_without_files(path: str) -> Generator[Path]:
+    """Yields every subdirectory of +path+ that has no files under it."""
+    return _subdirs_without_generic(path, lambda root, filenames: len(filenames) > 0)
 
-    directories = []
-    has_wheels = set()
 
-    for root_str, _, filenames in os.walk(Path(path).resolve()):
-        root = Path(root_str)
-        if any(x.endswith(".whl") for x in filenames):
-            # This directory contains a wheel file, so mark it and all of its
-            # parent directories as has-wheel.
-            # The last item in root.parents is ".", so we ignore it.
-            #
-            # Wrapping this in `list()` is only needed for Python 3.9.
-            has_wheels.update(list(root.parents)[:-1])
-            has_wheels.add(root)
-        directories.append(root)
-
-    for d in sorted(directories, reverse=True):
-        if d not in has_wheels:
-            yield d
+def subdirs_without_wheels(path: str) -> Generator[Path]:
+    """Yields every subdirectory of +path+ that has no .whl files under it."""
+    return _subdirs_without_generic(
+        path, lambda root, filenames: any(x.endswith(".whl") for x in filenames)
+    )
```

Thoughts?

Code after:
def directory_size(path: str) -> int | float:
    size = 0.0
    for root, _dirs, files in os.walk(path):
        for filename in files:
            file_path = os.path.join(root, filename)
            size += file_size(file_path)
    return size


def format_directory_size(path: str) -> str:
    return format_size(directory_size(path))


def copy_directory_permissions(directory: str, target_file: BinaryIO) -> None:
    mode = (
        os.stat(directory).st_mode & 0o666  # select read/write permissions of directory
        | 0o600  # set owner read/write permissions
    )
    # Change permissions only if there is no risk of following a symlink.
    if os.chmod in os.supports_fd:
        os.chmod(target_file.fileno(), mode)
    elif os.chmod in os.supports_follow_symlinks:
        os.chmod(target_file.name, mode, follow_symlinks=False)


def _subdirs_without_generic(
    path: str, predicate: Callable[[str, list[str]], bool]
) -> Generator[Path]:
    """Yields every subdirectory of +path+ that has no files matching the
    predicate under it."""

    directories = []
    excluded = set()

    for root_str, _, filenames in os.walk(Path(path).resolve()):
        root = Path(root_str)
        if predicate(root_str, filenames):
            # This directory should be excluded, so exclude it and all of its
            # parent directories.
            # The last item in root.parents is ".", so we ignore it.
            #
            # Wrapping this in `list()` is only needed for Python 3.9.
            excluded.update(list(root.parents)[:-1])
            excluded.add(root)
        directories.append(root)

    for d in sorted(directories, reverse=True):
        if d not in excluded:
            yield d