bug fix

Author: fangchenli

pandas/io/common.py

@@ -9,6 +9,7 @@
 import codecs
 from collections import defaultdict
 from collections.abc import (
+    Generator,
     Hashable,
     Mapping,
     Sequence,
@@ -26,7 +27,10 @@
 )
 import mmap
 import os
-from pathlib import Path
+from pathlib import (
+    Path,
+    PurePosixPath,
+)
 import re
 import tarfile
 from typing import (
@@ -1282,3 +1286,126 @@ def dedup_names(
         counts[col] = cur_count + 1
 
     return names
+
+
+def _match_file(
+    path: Path | PurePosixPath, extensions: set[str] | None, glob: str | None
+) -> bool:
+    """Check if the file matches the given extensions and glob pattern.
+    Parameters
+    ----------
+    path : Path or PurePosixPath
+        The file path to check.
+    extensions : set[str]
+        A set of file extensions to match against.
+    glob : str
+        A glob pattern to match against.
+    Returns
+    -------
+    bool
+        True if the file matches the extensions and glob pattern, False otherwise.
+    """
+    return (extensions is None or path.suffix.lower() in extensions) and (
+        glob is None or path.match(glob)
+    )
+
+
+def iterdir(
+    path: FilePath,
+    extensions: str | list[str] | None = None,
+    glob: str | None = None,
+) -> Generator[Path | PurePosixPath]:
+    """Yield file paths in a directory (no nesting allowed).
+
+    Supports:
+    - Local paths (str, os.PathLike)
+    - file:// URLs
+    - Remote paths (e.g., s3://) via fsspec (if installed)
+
+    Parameters
+    ----------
+    path : FilePath
+        Path to the directory (local or remote).
+    extensions : str or list of str, optional
+        Only yield files with the given extension(s). Case-insensitive.
+        If None, all files are yielded.
+    glob : str, optional
+        Only yield files matching the given glob pattern.
+        If None, all files are yielded.
+
+    Yields
+    ------
+    pathlib.Path or pathlib.PurePosixPath
+        File paths within the directory.
+
+    Raises
+    ------
+    NotADirectoryError
+        If the given path is not a directory.
+    ImportError
+        If fsspec is required but not installed.
+    """
+    if extensions is not None:
+        if isinstance(extensions, str):
+            extensions = {extensions.lower()}
+        else:
+            extensions = {ext.lower() for ext in extensions}
+
+    if isinstance(path, os.PathLike):
+        path = os.fspath(path)
+
+    parsed = parse_url(path)
+    scheme = parsed.scheme or "file"
+    base_path = parsed.path if scheme == "file" else path
+
+    if scheme == "file":
+        resolved_path = Path(base_path)
+        if resolved_path.is_file():
+            if _match_file(
+                resolved_path,
+                extensions,
+                glob,
+            ):
+                yield resolved_path
+            return
+
+        if not resolved_path.is_dir():
+            raise NotADirectoryError(
+                f"Path {path!r} is neither a file nor a directory."
+            )
+
+        for entry in resolved_path.iterdir():
+            if entry.is_file():
+                if _match_file(
+                    entry,
+                    extensions,
+                    glob,
+                ):
+                    yield entry
+        return
+
+    # Remote paths (e.g., s3)
+    fsspec = import_optional_dependency("fsspec", extra=scheme)
+    fs = fsspec.filesystem(scheme)
+    if fs.isfile(path):
+        path_obj = PurePosixPath(path)
+        if _match_file(
+            path_obj,
+            extensions,
+            glob,
+        ):
+            yield path_obj
+        return
+    if not fs.isdir(path):
+        raise NotADirectoryError(f"Remote path {path!r} is not a directory.")
+
+    files = fs.ls(path, detail=True)
+    for f in files:
+        if f["type"] == "file":
+            path_obj = PurePosixPath(f["name"])
+            if _match_file(
+                path_obj,
+                extensions,
+                glob,
+            ):
+                yield path_obj

pandas/tests/io/conftest.py

@@ -223,3 +223,14 @@ def compression_format(request):
 @pytest.fixture(params=_compression_formats_params)
 def compression_ext(request):
     return request.param[0]
+
+
+@pytest.fixture
+def directory_with_dummy_csv(tmp_path):
+    """
+    Fixture to create a directory with dummy CSV files for testing.
+    """
+    for i in range(3):
+        file_path = tmp_path / f"file_{i}.csv"
+        file_path.touch()
+    return tmp_path

pandas/tests/io/parser/test_directory.py

@@ -0,0 +1,38 @@
+from csv import (
+    DictWriter,
+    reader as csv_reader,
+)
+
+import pytest
+
+
+@pytest.fixture
+def directory_data():
+    return ["a", "b", "c"], [
+        {"first": {"a": 1, "b": 2, "c": 3}},
+        {"second": {"a": 4, "b": 5, "c": 6}},
+        {"third": {"a": 7, "b": 8, "c": 9}},
+    ]
+
+
+@pytest.fixture
+def directory_data_to_file(tmp_path, directory_data):
+    field_names, data_list = directory_data
+    for data in data_list:
+        file_name = next(iter(data.keys()))
+        path = tmp_path / f"{file_name}.csv"
+        with path.open("w", newline="") as file:
+            writer = DictWriter(file, fieldnames=field_names)
+            writer.writeheader()
+            writer.writerow(data[file_name])
+    return tmp_path
+
+
+def test_directory_data(directory_data_to_file):
+    print(type(directory_data_to_file.iterdir()))
+    assert len(list(directory_data_to_file.iterdir())) == 3
+    for file in directory_data_to_file.iterdir():
+        with file.open() as f:
+            reader = csv_reader(f)
+            header = next(reader)
+            assert header == ["a", "b", "c"]

pandas/tests/io/test_common.py

@@ -695,3 +695,10 @@ def test_pyarrow_read_csv_datetime_dtype():
     expect = pd.DataFrame({"date": expect_data})
 
     tm.assert_frame_equal(expect, result)
+
+
+def test_iterdir(directory_with_dummy_csv):
+    for file in icom.iterdir(directory_with_dummy_csv):
+        assert file.is_file()
+        assert file.name.startswith("file_")
+        assert file.suffix == ".csv"