mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-07 23:09:21 +01:00
utils.http: remove open url (#6700)
This was used before on opening http urls as a stream, but we no longer use those (as we use fsspec's open method). One thing to note however is that, compared to previous implementation, fsspec's http filesystem's open is not very robust to failures.
This commit is contained in:
@@ -1,67 +0,0 @@
|
||||
import io
|
||||
from contextlib import contextmanager
|
||||
|
||||
from dvc.utils.stream import IterStream
|
||||
|
||||
|
||||
@contextmanager
|
||||
def open_url(url, mode="r", encoding=None, **iter_opts):
|
||||
"""Opens an url as a readable stream.
|
||||
|
||||
Resumes on connection error.
|
||||
Url could be a string or a callable returning a string.
|
||||
"""
|
||||
assert mode in {"r", "rt", "rb"}
|
||||
|
||||
with iter_url(url, **iter_opts) as (response, it):
|
||||
bytes_stream = IterStream(it)
|
||||
|
||||
if mode == "rb":
|
||||
yield bytes_stream
|
||||
else:
|
||||
encoding = encoding or response.encoding
|
||||
yield io.TextIOWrapper(bytes_stream, encoding=encoding)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def iter_url(url, chunk_size=io.DEFAULT_BUFFER_SIZE, auth=None):
|
||||
"""Iterate over chunks requested from url."""
|
||||
import requests
|
||||
|
||||
def request(headers=None):
|
||||
the_url = url() if callable(url) else url
|
||||
response = requests.get(
|
||||
the_url, stream=True, headers=headers, auth=auth
|
||||
)
|
||||
if response.status_code == 404:
|
||||
raise FileNotFoundError(f"Can't open {the_url}")
|
||||
response.raise_for_status()
|
||||
return response
|
||||
|
||||
def gen(response):
|
||||
try:
|
||||
pos = 0
|
||||
while True:
|
||||
try:
|
||||
for chunk in response.iter_content(chunk_size):
|
||||
pos += len(chunk)
|
||||
yield chunk
|
||||
break
|
||||
except requests.ConnectionError:
|
||||
response.close()
|
||||
if response.headers.get("Accept-Ranges") != "bytes":
|
||||
raise
|
||||
|
||||
# Reopen request from where we stopped
|
||||
headers = {"Range": f"bytes={pos}-"}
|
||||
response = request(headers)
|
||||
finally:
|
||||
response.close()
|
||||
|
||||
response = request()
|
||||
it = gen(response)
|
||||
try:
|
||||
yield response, it
|
||||
finally:
|
||||
# Ensure connection is closed
|
||||
it.close()
|
||||
@@ -8,61 +8,6 @@ from dvc.istextfile import DEFAULT_CHUNK_SIZE, istextblock
|
||||
from dvc.utils import dos2unix
|
||||
|
||||
|
||||
class IterStream(io.RawIOBase):
|
||||
"""Wraps an iterator yielding bytes as a file object"""
|
||||
|
||||
def __init__(self, iterator): # pylint: disable=super-init-not-called
|
||||
self.iterator = iterator
|
||||
self.leftover = b""
|
||||
|
||||
def readable(self):
|
||||
return True
|
||||
|
||||
def writable(self) -> bool:
|
||||
return False
|
||||
|
||||
# Python 3 requires only .readinto() method, it still uses other ones
|
||||
# under some circumstances and falls back if those are absent. Since
|
||||
# iterator already constructs byte strings for us, .readinto() is not the
|
||||
# most optimal, so we provide .read1() too.
|
||||
|
||||
def readinto(self, b):
|
||||
try:
|
||||
n = len(b) # We're supposed to return at most this much
|
||||
chunk = self.leftover or next(self.iterator)
|
||||
output, self.leftover = chunk[:n], chunk[n:]
|
||||
|
||||
n_out = len(output)
|
||||
b[:n_out] = output
|
||||
return n_out
|
||||
except StopIteration:
|
||||
return 0 # indicate EOF
|
||||
|
||||
readinto1 = readinto
|
||||
|
||||
def read1(self, n=-1):
|
||||
try:
|
||||
chunk = self.leftover or next(self.iterator)
|
||||
except StopIteration:
|
||||
return b""
|
||||
|
||||
# Return an arbitrary number or bytes
|
||||
if n <= 0:
|
||||
self.leftover = b""
|
||||
return chunk
|
||||
|
||||
output, self.leftover = chunk[:n], chunk[n:]
|
||||
return output
|
||||
|
||||
def peek(self, n):
|
||||
while len(self.leftover) < n:
|
||||
try:
|
||||
self.leftover += next(self.iterator)
|
||||
except StopIteration:
|
||||
break
|
||||
return self.leftover[:n]
|
||||
|
||||
|
||||
class HashedStreamReader(io.IOBase):
|
||||
|
||||
PARAM_CHECKSUM = "md5"
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
import io
|
||||
|
||||
import requests
|
||||
|
||||
from dvc.utils.http import open_url
|
||||
|
||||
|
||||
def read_nbytes(fd, num, mode="r"):
|
||||
data = b"" if mode == "rb" else ""
|
||||
while len(data) < num:
|
||||
chunk = fd.read(num - len(data))
|
||||
if not chunk:
|
||||
break
|
||||
data += chunk
|
||||
return data
|
||||
|
||||
|
||||
def test_open_url(tmp_path, monkeypatch, http):
|
||||
# Simulate bad connection
|
||||
original_iter_content = requests.Response.iter_content
|
||||
|
||||
def bad_iter_content(self, *args, **kwargs):
|
||||
it = original_iter_content(self, *args, **kwargs)
|
||||
for i, chunk in enumerate(it):
|
||||
# Drop connection error on second chunk if there is one
|
||||
if i > 0:
|
||||
raise requests.ConnectionError("Simulated connection drop")
|
||||
yield chunk
|
||||
|
||||
monkeypatch.setattr(requests.Response, "iter_content", bad_iter_content)
|
||||
|
||||
# Text should be longer than default chunk to test resume,
|
||||
# using twice of that plus something tests second resume,
|
||||
# this is important because second response is different
|
||||
text = "0123456789" * (io.DEFAULT_BUFFER_SIZE // 10 + 1)
|
||||
http.gen("sample.txt", text * 2)
|
||||
|
||||
with open_url((http / "sample.txt").url) as fd:
|
||||
# Test various .read() variants
|
||||
assert read_nbytes(fd, len(text), mode="r") == text
|
||||
assert fd.read() == text
|
||||
assert fd.read() == ""
|
||||
|
||||
|
||||
def test_open_url_peek_rb(tmp_path, monkeypatch, http):
|
||||
# Goes over seek feature in 'rb' mode
|
||||
text = "0123456789" * (io.DEFAULT_BUFFER_SIZE // 10 + 1)
|
||||
http.gen("sample.txt", text * 2)
|
||||
|
||||
with open_url((http / "sample.txt").url, mode="rb") as fd:
|
||||
text = text.encode("utf8")
|
||||
assert fd.peek(len(text)) == text # pylint: disable=no-member
|
||||
assert read_nbytes(fd, len(text), mode="rb") == text
|
||||
assert fd.peek(len(text)) == text # pylint: disable=no-member
|
||||
assert fd.read() == text
|
||||
assert fd.read() == b""
|
||||
Reference in New Issue
Block a user