Skip to content

graflo.data_source.files

Files read through a files connection: a local directory or an S3 prefix.

A FileConnector bound to a :class:~graflo.connections.sources.FilesConnConfig reads sub_path relative to the connection's root. :func:list_files finds the files it matches; :func:file_data_source reads one of them. Both serve sampling and ingestion alike.

Attributes

__all__ = ['FileEntry', 'S3ObjectDataSource', 'file_data_source', 'list_files', 'walk'] module-attribute

Classes

FileEntry

One file under a files connection: its path relative to the root, and size.

Source code in graflo/data_source/files.py 
28
29
30
31
32
33
34
35
36
37
38
class FileEntry:
    """One file under a files connection: its path relative to the root, and size."""

    __slots__ = ("relative", "size")

    def __init__(self, relative: str, size: int | None) -> None:
        self.relative = relative
        self.size = size

    def __repr__(self) -> str:
        return f"FileEntry({self.relative!r}, {self.size!r})"

Attributes

__slots__ = ('relative', 'size') class-attribute instance-attribute
relative = relative instance-attribute
size = size instance-attribute

Methods:

__init__(relative, size)
Source code in graflo/data_source/files.py 
33
34
35
def __init__(self, relative: str, size: int | None) -> None:
    self.relative = relative
    self.size = size
__repr__()
Source code in graflo/data_source/files.py 
37
38
def __repr__(self) -> str:
    return f"FileEntry({self.relative!r}, {self.size!r})"

S3ObjectDataSource

Bases: AbstractDataSource

One S3 object, read as the file it is: downloaded, then read like a local file.

The download lives only while :meth:iter_batches runs.

Source code in graflo/data_source/files.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
class S3ObjectDataSource(AbstractDataSource):
    """One S3 object, read as the file it is: downloaded, then read like a local file.

    The download lives only while :meth:`iter_batches` runs.
    """

    s3: S3ConnConfig
    key: str
    file_type: str | None = None
    encoding: EncodingType = EncodingType.UTF_8
    source_type: DataSourceType = DataSourceType.FILE
    _client: Any = PrivateAttr(default=None)

    def iter_batches(
        self, batch_size: int = 1000, limit: int | None = None
    ) -> Iterator[list[dict]]:
        from graflo.object_storage.s3_client import boto3_s3_client

        client = self._client or boto3_s3_client(self.s3)
        name = PurePosixPath(self.key).name
        with tempfile.TemporaryDirectory(prefix="graflo-s3-") as tmp:
            local = Path(tmp) / name
            client.download_file(self.s3.bucket, self.key, str(local))
            source = DataSourceFactory.create_file_data_source(
                path=local, file_type=self.file_type, encoding=self.encoding
            )
            yield from source.iter_batches(batch_size=batch_size, limit=limit)

Attributes

encoding = EncodingType.UTF_8 class-attribute instance-attribute
file_type = None class-attribute instance-attribute
key instance-attribute
s3 instance-attribute
source_type = DataSourceType.FILE class-attribute instance-attribute

Methods:

iter_batches(batch_size=1000, limit=None)
Source code in graflo/data_source/files.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def iter_batches(
    self, batch_size: int = 1000, limit: int | None = None
) -> Iterator[list[dict]]:
    from graflo.object_storage.s3_client import boto3_s3_client

    client = self._client or boto3_s3_client(self.s3)
    name = PurePosixPath(self.key).name
    with tempfile.TemporaryDirectory(prefix="graflo-s3-") as tmp:
        local = Path(tmp) / name
        client.download_file(self.s3.bucket, self.key, str(local))
        source = DataSourceFactory.create_file_data_source(
            path=local, file_type=self.file_type, encoding=self.encoding
        )
        yield from source.iter_batches(batch_size=batch_size, limit=limit)

Functions:

file_data_source(files, relative, *, file_type=None, encoding=EncodingType.UTF_8)

The data source reading the file at relative under the root of files.

Source code in graflo/data_source/files.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def file_data_source(
    files: FilesConnConfig,
    relative: str,
    *,
    file_type: str | None = None,
    encoding: EncodingType = EncodingType.UTF_8,
) -> AbstractDataSource:
    """The data source reading the file at *relative* under the root of *files*."""
    if files.is_s3:
        return S3ObjectDataSource(
            s3=files.s3_config(),
            key=_s3_key(files, relative),
            file_type=file_type,
            encoding=encoding,
        )
    return DataSourceFactory.create_file_data_source(
        path=files.local_root / relative, file_type=file_type, encoding=encoding
    )

list_files(files, connector, limit=None)

Paths, relative to the root of files, of the files connector matches.

connector.regex matches a file's name, as for a local directory.

Source code in graflo/data_source/files.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def list_files(
    files: FilesConnConfig, connector: FileConnector, limit: int | None = None
) -> list[str]:
    """Paths, relative to the root of *files*, of the files *connector* matches.

    ``connector.regex`` matches a file's name, as for a local directory.
    """
    entries = walk(files, str(connector.sub_path), recursive=connector.recursive)
    matched = [
        entry.relative
        for entry in entries
        if connector.regex is None
        or re.search(connector.regex, PurePosixPath(entry.relative).name) is not None
    ]
    return matched[:limit] if limit is not None else matched

walk(files, sub_path='', *, recursive=True, limit=None)

Every file under sub_path of files, sorted by relative path.

Parameters:

Name Type Description Default
files FilesConnConfig

The files connection.

 required
sub_path str

Directory relative to the connection's root.

 ''
recursive bool

Also list the directories below sub_path.

 True
limit int | None

Stop after this many files.

 None
Source code in graflo/data_source/files.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def walk(
    files: FilesConnConfig,
    sub_path: str = "",
    *,
    recursive: bool = True,
    limit: int | None = None,
) -> list[FileEntry]:
    """Every file under *sub_path* of *files*, sorted by relative path.

    Args:
        files: The files connection.
        sub_path: Directory relative to the connection's root.
        recursive: Also list the directories below *sub_path*.
        limit: Stop after this many files.
    """
    base = str(PurePosixPath(sub_path)).strip("/")
    base = "" if base == "." else base
    found: list[FileEntry] = []
    if files.is_s3:
        client = _s3_client(files)
        root_key = _s3_key(files, base)
        prefix = f"{root_key}/" if root_key else ""
        kwargs: dict[str, Any] = {"Bucket": files.bucket, "Prefix": prefix}
        if not recursive:
            kwargs["Delimiter"] = "/"
        strip = f"{files.prefix}/" if files.prefix else ""
        for page in client.get_paginator("list_objects_v2").paginate(**kwargs):
            for item in page.get("Contents", []):
                key = str(item["Key"])
                if key.endswith("/"):
                    continue
                found.append(FileEntry(key[len(strip) :], item.get("Size")))
    else:
        root = files.local_root
        directory = root / base if base else root
        if directory.is_dir():
            candidates = directory.rglob("*") if recursive else directory.iterdir()
            found = [
                FileEntry(path.relative_to(root).as_posix(), path.stat().st_size)
                for path in candidates
                if path.is_file()
            ]
    found.sort(key=lambda entry: entry.relative)
    return found[:limit] if limit is not None else found