File size: 5,205 Bytes
6d63e5b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
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
141
from typing import TYPE_CHECKING, Optional
from urllib.parse import urlparse

from ._cached_path import cached_path, get_from_cache
from .common import PathOrStr
from .schemes import get_scheme_client, get_supported_schemes

if TYPE_CHECKING:
    from rich.progress import Progress


def get_bytes_range(
    url_or_filename: PathOrStr,
    index: int,
    length: int,
    cache_dir: Optional[PathOrStr] = None,
    extract_archive: bool = False,
    force_extract: bool = False,
    quiet: bool = False,
    progress: Optional["Progress"] = None,
) -> bytes:
    """
    Get a range of up to ``length`` bytes starting at ``index``.

    In some cases the entire file may need to be downloaded, such as when the server does not support
    a range download or when you're trying to get a bytes range from a file within an archive.

    .. caution::
        You may get less than ``length`` bytes sometimes, such as when fetching a range from an HTTP
        resource starting at 0 since headers will be omitted in the bytes returned.

    Parameters
    ----------

    url_or_filename :
        A URL or path to parse and possibly download.

    index :
        The index of the byte to start at.

    length :
        The number of bytes to read.

    cache_dir :
        The directory to cache downloads. If not specified, the global default cache directory
        will be used (``~/.cache/cached_path``). This can be set to something else with
        :func:`set_cache_dir()`.

        This is only relevant when the bytes range cannot be obtained directly from the resource.

    extract_archive :
        Set this to ``True`` when you want to get a bytes range from a file within an archive.
        In this case the ``url_or_filename`` must contain an "!" followed by the relative path of the file
        within the archive, e.g. "s3://my-archive.tar.gz!my-file.txt".

        Note that the entire archive has to be downloaded in this case.

    force_extract :
        If ``True`` and the resource is a file within an archive (when the path contains an "!" and
        ``extract_archive=True``), it will be extracted regardless of whether or not the extracted
        directory already exists.

        .. caution::
            Use this flag with caution! This can lead to race conditions if used
            from multiple processes on the same file.

    quiet :
        If ``True``, progress displays won't be printed.

        This is only relevant when the bytes range cannot be obtained directly from the resource.

    progress :
        A custom progress display to use. If not set and ``quiet=False``, a default display
        from :func:`~cached_path.get_download_progress()` will be used.

        This is only relevant when the bytes range cannot be obtained directly from the resource.
    """
    if not isinstance(url_or_filename, str):
        url_or_filename = str(url_or_filename)

    # If we're using the /a/b/foo.zip!c/d/file.txt syntax, handle it here.
    exclamation_index = url_or_filename.find("!")
    if extract_archive and exclamation_index >= 0:
        archive_path = url_or_filename[:exclamation_index]
        file_name = url_or_filename[exclamation_index + 1 :]

        # Call 'cached_path' now to get the local path to the archive itself.
        cached_archive_path = cached_path(
            archive_path,
            cache_dir=cache_dir,
            extract_archive=True,
            force_extract=force_extract,
            quiet=quiet,
            progress=progress,
        )
        if not cached_archive_path.is_dir():
            raise ValueError(
                f"{url_or_filename} uses the ! syntax, but does not specify an archive file."
            )

        # Now load bytes from the desired file within the extracted archive, provided it exists.
        file_path = cached_archive_path / file_name
        if not file_path.exists():
            raise FileNotFoundError(f"'{file_name}' not found within '{archive_path}'")

        return _bytes_range_from_file(file_path, index, length)

    if urlparse(url_or_filename).scheme in get_supported_schemes():
        # URL, so use the scheme client.
        client = get_scheme_client(url_or_filename)

        # Check if file is already downloaded.
        try:
            cache_path, _ = get_from_cache(
                url_or_filename,
                cache_dir=cache_dir,
                quiet=quiet,
                progress=progress,
                no_downloads=True,
                _client=client,
            )
            return _bytes_range_from_file(cache_path, index, length)
        except FileNotFoundError:
            pass

        # Otherwise try streaming bytes directly.
        try:
            return client.get_bytes_range(index, length)
        except NotImplementedError:
            # fall back to downloading the whole file.
            pass

    file_path = cached_path(url_or_filename, cache_dir=cache_dir, quiet=quiet, progress=progress)
    return _bytes_range_from_file(file_path, index, length)


def _bytes_range_from_file(path: PathOrStr, index: int, length: int) -> bytes:
    with open(path, "rb") as f:
        f.seek(index)
        return f.read(length)