From c64b4ce78f553c40a7916e1857bbdd63ff7891ab Mon Sep 17 00:00:00 2001
From: Dmitry Kropachev <40304587+dkropachev@users.noreply.github.com>
Date: Tue, 29 Dec 2020 04:34:24 +0700
Subject: [PATCH] dvc/utils/stream: make IterStream seekable and peekable
 (#5084)

fixes #5021
---
 dvc/utils/stream.py           |  15 +++++++++++++--
 tests/test.p12                | Bin 0 -> 2479 bytes
 tests/unit/remote/test.p12    | Bin 0 -> 2479 bytes
 tests/unit/utils/test_http.py |  27 ++++++++++++++++++++++++++-
 4 files changed, 39 insertions(+), 3 deletions(-)
 create mode 100644 tests/test.p12
 create mode 100644 tests/unit/remote/test.p12

diff --git a/dvc/utils/stream.py b/dvc/utils/stream.py
index cfa18d183..c2d3fe33e 100644
--- a/dvc/utils/stream.py
+++ b/dvc/utils/stream.py
@@ -6,11 +6,14 @@ class IterStream(io.RawIOBase):
 
     def __init__(self, iterator):  # pylint: disable=super-init-not-called
         self.iterator = iterator
-        self.leftover = None
+        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
@@ -38,8 +41,16 @@ class IterStream(io.RawIOBase):
 
         # Return an arbitrary number or bytes
         if n <= 0:
-            self.leftover = None
+            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]
diff --git a/tests/test.p12 b/tests/test.p12
new file mode 100644
index 0000000000000000000000000000000000000000..6c391b054f9d0551b53682c190206c6e294fca25
GIT binary patch
literal 2479
zcmY+EXEYlM8^=S4+Iv)u64ls2irQQ*Mbu1V#2yW*gj%IEN>zzHi_%^y_9$AaMoR4(
zEh#l?$4+o-UhjL}_ug|qJm)#*_y0f7#~&0+UrtR0gktI4K=hKahOvjtRCH8@SbA$3
zEWP;!hC{Km{{M<-U1+eh4j0(=Vh|w4|6MULP*E3RX^x>-nnS1rh~fYE;`tRCuJ;^F
z^RR+co&)*!Wl0MpZhu5?IW-VKFr&fJ^c-!_#@CIl%>`GKTYp_g)E@?=zE&Cg!ozmW
z)QEIH>VVuDryWzT<Z-(2RceR`(|>Kocp&_)FmtDo7%LPQz0vALxtlk?|1Dm}wD3fj
zD?6&PBSAwpq5fKs7i!ayh7xD+#+92xw)l#*8?(zJuei{+OaV$3K04iMiG1aTZ}lHy
z{bi{ADPRv00+odEpN{|wL9#UkHDP?)bo`X24c}U#K4NKbq{Qbb*_Mt@UdxTO<eg92
zVHm_XaF16miS+CAnDhCUJci7-6`IJBGLp-U@4+S<TM{|UPXe2ZDm|BQWKq`B^4^x2
zM{SHMkC$<`m!uVX^?2lh9ArDl2F8ttUWYaEYC$F2B4doWN^AsgTAPn|dPa7gLr<xO
zhXF3*_p%^fKb!jVRrDWXTZ&xLNjQGFXV`#2t2q0hSDidG&8jwjruz?Ex}Xf$e0#Iq
zlw^E>k;tV^wVb1er<@4ih_iyrWv=!=6bjPv-lPwcj!*hSGSgdYOxN>0&-{&NI-`1!
zu+mYvp-KJ0Ys|XI)%|@QRQopAdGaSjj3i2;c;vD)T1i?Zr@{EF6xkOpw%rb4`^B@u
zDscz7P1rRtcAI?wf8gshYAJo%*8`_JM$QE&rk?tu9cQ|gXSkc?qS*a8-LB&ooEwO>
zNW_uT-fHkR@0Fz81O0^3e3m)10njztc=x><=uUNh)l!-2<Zp5O9}`m0f{ZVorb@nW
z4bYeCKaCz%;M*ylc!$SPF7H39h@OF+joH*c9@vhAj`+VB6;RZczw{o}b`rS0!q%gV
zu1!mnK!!`bU!4-z%E;BgufT^zlLkP2Zr<tf_M_f)lhc3@JzTZP!K!7Mh662&*ioiM
zgup8NSHHYYtX)7B-cDW6?K`XgMCZ%i6){ZHc-zeyd~T&|GV=y<Mu|_jTj38Vh049?
zln<fV)tBpZ@0D|8029&GhI^6}(ZBO&9J@DaY@JE@XSZF~0@__$26n<dLwNd%G*Q`g
zwZob;`gt{saR9zJ{F!(jKjyF7nxF8s8*FVt%cpxy`06HvVBIeLNT$|)-%u_PoQ<%e
z<agFUd~i2p^A2NXlb2rh;=PE&Be3Rv4r8j+hIBBu+46bwv!k#dpa-Fd8(O;p9YF&U
z!Xo~<IuoKE8m0%d%eU(?f+Is;#4C2Eddj|Bt97RC^HWFOj`WHAA!r5e-$aF^wv&D9
zHY)QgM=;aRmqaOte>(Oj04*C0b<cm7d35LOG94X(nIJ{N?2xu9Q>TwBFjMc$p*p0Y
zTd_tJ$}R+ThWHQO(XUn$Obm<$#`Q!dE$TYD%ZtsJZRdoI7|PD0Ld#6`qDzux#l2+@
zmLpe!jQKN8Z9+}F8G}{!w3CdFq?cO?M;X2dZmhjQxmIx+4MZtP7N8@<Pf*~G-#!dz
z3FyYdiwS+&x^OFtiY8+{HKwFgTBWR~w!U}*`ovxCs$8Ntc{}e9RNXIj1%X{<=SR0C
zJ#5=pwnNPJ!VVbIpeZR<X#L=UTt^{!c{*(w*}QqqHNl6(gNW#t+HdZ9_NnHLW3?4u
zXo6S~dZ%%gPJGEO7%{?*F^4BCmmX6n>v0A|?!nt#HB|47Y-WzT_fJWB7z-X1h;+*v
zLcxEllAcSRM~DmH2fzT(fFOV~AQ0dV@V=N(sK`GhZh0;m4lDFM_p9<sP!%Xl`6l!x
z6s7{j0>A!Uq5~CTfn67<gPIC(v6}x$sQzVI(0@!z_4MlbvV~m~kRSGxQxSJG)4$gA
zU(-fmfl=;-6POB?Fvxc`czu_T;y(9r!?j&x){xE4-IidW0ovZ#@-Yp|%Thr4;RXy5
zHrndss~$?#5HMZ*DrKf4SFITSOj$F_&La!z!hf@HiVA(S8SA>O2Mijq!vAQ#{~6I5
zm-epoa|Zkncmui8@-a6YzPZEqX6mDAvA7|_+ew9UAz}9xK68ciuh)8kD@9-YP`&Bj
zRr{b>3Z8d&bMe5hHuoR)JkHV?dfr+y&bF3rZ+0`GNV>t(nPW^4@X5zwawj|{Ov+)_
zF{p9-Qi|ob^(y^)pX@A4^n)!;bv*YNE9J%Dht?(=3lP@~+P!dIl0_+3`)m@@!bYd1
zB<9}X81S39-r#3-$i4YAP+-|Ou3@!6TPnm9tgezZ$;Mk|E_(HF-kn1bMXfclP*bC9
zv`_6k2>Kv#YSyx9Ka3lGfj%ogHv(d&@ypCb5KI={y~=I%PUlf>r#S3(=ajy-8Ogg;
zG28F`E=XMCRX7NImN$JF)x?(FX7VfH#!}@eqz}_$gvgvX;1{BnHE6DYc<0D<Qv2FX
zs69*B%K#5tZ(XeJC9f|#j9f3?XxAM&gEV2uY8=Jm;)`O}OCsxUoe*7%p_6@kgUsJ=
zxdTFUBBU^6XvDz*(<jm0W=2QR-&=hd12O1?PmM=P0nKuaEIbpx5%%voa>Pkeh~XMv
zWM^#1Vgd`yS6E7Xl2~+Aow+`2kzILmLb-(0VL_#K4==^A4pI8_63HzxepN&T3jic;
zgS~z>C3J)FNp)f?-N)?Jz}+recJ4nnEm7<`Do?Z=eDb6STeCmW5RgsFZmFh{99V34
z8aza3kqVDJS;;9|A76u{&Lw%ce+kLv#<`mVH0(5&oZ53bWyTYRnAC@J@OUu%bsWyu
z!*|3<t^oQQTgjegOHG%lyug@IW|08iXrihpM+s$09YSxWxR?Mu*Gm5Njpy_7s^wg^
zH`K|DO!;|i*#6v=$sD1dd*f$8Fxb*<LE|~%2XW0OJLXd>v5B|o0r+t%?S9p_#;qwf
z1ucjj=Vfj$?r66RY{h+Z$<*zUCzkmLF$vR?e4g1lF{DK^`*%);+2mV<#d%r#9K%72
zs%m<-j}(`wL~m)iA+=&k=kOw>@HKMTA>h;*Xp@r?JZ`jSgLf<IwH2(**tJEpsMpRf
zC91F)?w8IGeB$xS<BNt+O{fTz9z-j}0;J}r1<-JXRR8o6E9~ugbmiS+K@~WBla&%e
f!$qI9=&$lp6&r<XJDQs#j!ueBL{bB(VhjHQoyMDq

literal 0
HcmV?d00001

diff --git a/tests/unit/remote/test.p12 b/tests/unit/remote/test.p12
new file mode 100644
index 0000000000000000000000000000000000000000..6c391b054f9d0551b53682c190206c6e294fca25
GIT binary patch
literal 2479
zcmY+EXEYlM8^=S4+Iv)u64ls2irQQ*Mbu1V#2yW*gj%IEN>zzHi_%^y_9$AaMoR4(
zEh#l?$4+o-UhjL}_ug|qJm)#*_y0f7#~&0+UrtR0gktI4K=hKahOvjtRCH8@SbA$3
zEWP;!hC{Km{{M<-U1+eh4j0(=Vh|w4|6MULP*E3RX^x>-nnS1rh~fYE;`tRCuJ;^F
z^RR+co&)*!Wl0MpZhu5?IW-VKFr&fJ^c-!_#@CIl%>`GKTYp_g)E@?=zE&Cg!ozmW
z)QEIH>VVuDryWzT<Z-(2RceR`(|>Kocp&_)FmtDo7%LPQz0vALxtlk?|1Dm}wD3fj
zD?6&PBSAwpq5fKs7i!ayh7xD+#+92xw)l#*8?(zJuei{+OaV$3K04iMiG1aTZ}lHy
z{bi{ADPRv00+odEpN{|wL9#UkHDP?)bo`X24c}U#K4NKbq{Qbb*_Mt@UdxTO<eg92
zVHm_XaF16miS+CAnDhCUJci7-6`IJBGLp-U@4+S<TM{|UPXe2ZDm|BQWKq`B^4^x2
zM{SHMkC$<`m!uVX^?2lh9ArDl2F8ttUWYaEYC$F2B4doWN^AsgTAPn|dPa7gLr<xO
zhXF3*_p%^fKb!jVRrDWXTZ&xLNjQGFXV`#2t2q0hSDidG&8jwjruz?Ex}Xf$e0#Iq
zlw^E>k;tV^wVb1er<@4ih_iyrWv=!=6bjPv-lPwcj!*hSGSgdYOxN>0&-{&NI-`1!
zu+mYvp-KJ0Ys|XI)%|@QRQopAdGaSjj3i2;c;vD)T1i?Zr@{EF6xkOpw%rb4`^B@u
zDscz7P1rRtcAI?wf8gshYAJo%*8`_JM$QE&rk?tu9cQ|gXSkc?qS*a8-LB&ooEwO>
zNW_uT-fHkR@0Fz81O0^3e3m)10njztc=x><=uUNh)l!-2<Zp5O9}`m0f{ZVorb@nW
z4bYeCKaCz%;M*ylc!$SPF7H39h@OF+joH*c9@vhAj`+VB6;RZczw{o}b`rS0!q%gV
zu1!mnK!!`bU!4-z%E;BgufT^zlLkP2Zr<tf_M_f)lhc3@JzTZP!K!7Mh662&*ioiM
zgup8NSHHYYtX)7B-cDW6?K`XgMCZ%i6){ZHc-zeyd~T&|GV=y<Mu|_jTj38Vh049?
zln<fV)tBpZ@0D|8029&GhI^6}(ZBO&9J@DaY@JE@XSZF~0@__$26n<dLwNd%G*Q`g
zwZob;`gt{saR9zJ{F!(jKjyF7nxF8s8*FVt%cpxy`06HvVBIeLNT$|)-%u_PoQ<%e
z<agFUd~i2p^A2NXlb2rh;=PE&Be3Rv4r8j+hIBBu+46bwv!k#dpa-Fd8(O;p9YF&U
z!Xo~<IuoKE8m0%d%eU(?f+Is;#4C2Eddj|Bt97RC^HWFOj`WHAA!r5e-$aF^wv&D9
zHY)QgM=;aRmqaOte>(Oj04*C0b<cm7d35LOG94X(nIJ{N?2xu9Q>TwBFjMc$p*p0Y
zTd_tJ$}R+ThWHQO(XUn$Obm<$#`Q!dE$TYD%ZtsJZRdoI7|PD0Ld#6`qDzux#l2+@
zmLpe!jQKN8Z9+}F8G}{!w3CdFq?cO?M;X2dZmhjQxmIx+4MZtP7N8@<Pf*~G-#!dz
z3FyYdiwS+&x^OFtiY8+{HKwFgTBWR~w!U}*`ovxCs$8Ntc{}e9RNXIj1%X{<=SR0C
zJ#5=pwnNPJ!VVbIpeZR<X#L=UTt^{!c{*(w*}QqqHNl6(gNW#t+HdZ9_NnHLW3?4u
zXo6S~dZ%%gPJGEO7%{?*F^4BCmmX6n>v0A|?!nt#HB|47Y-WzT_fJWB7z-X1h;+*v
zLcxEllAcSRM~DmH2fzT(fFOV~AQ0dV@V=N(sK`GhZh0;m4lDFM_p9<sP!%Xl`6l!x
z6s7{j0>A!Uq5~CTfn67<gPIC(v6}x$sQzVI(0@!z_4MlbvV~m~kRSGxQxSJG)4$gA
zU(-fmfl=;-6POB?Fvxc`czu_T;y(9r!?j&x){xE4-IidW0ovZ#@-Yp|%Thr4;RXy5
zHrndss~$?#5HMZ*DrKf4SFITSOj$F_&La!z!hf@HiVA(S8SA>O2Mijq!vAQ#{~6I5
zm-epoa|Zkncmui8@-a6YzPZEqX6mDAvA7|_+ew9UAz}9xK68ciuh)8kD@9-YP`&Bj
zRr{b>3Z8d&bMe5hHuoR)JkHV?dfr+y&bF3rZ+0`GNV>t(nPW^4@X5zwawj|{Ov+)_
zF{p9-Qi|ob^(y^)pX@A4^n)!;bv*YNE9J%Dht?(=3lP@~+P!dIl0_+3`)m@@!bYd1
zB<9}X81S39-r#3-$i4YAP+-|Ou3@!6TPnm9tgezZ$;Mk|E_(HF-kn1bMXfclP*bC9
zv`_6k2>Kv#YSyx9Ka3lGfj%ogHv(d&@ypCb5KI={y~=I%PUlf>r#S3(=ajy-8Ogg;
zG28F`E=XMCRX7NImN$JF)x?(FX7VfH#!}@eqz}_$gvgvX;1{BnHE6DYc<0D<Qv2FX
zs69*B%K#5tZ(XeJC9f|#j9f3?XxAM&gEV2uY8=Jm;)`O}OCsxUoe*7%p_6@kgUsJ=
zxdTFUBBU^6XvDz*(<jm0W=2QR-&=hd12O1?PmM=P0nKuaEIbpx5%%voa>Pkeh~XMv
zWM^#1Vgd`yS6E7Xl2~+Aow+`2kzILmLb-(0VL_#K4==^A4pI8_63HzxepN&T3jic;
zgS~z>C3J)FNp)f?-N)?Jz}+recJ4nnEm7<`Do?Z=eDb6STeCmW5RgsFZmFh{99V34
z8aza3kqVDJS;;9|A76u{&Lw%ce+kLv#<`mVH0(5&oZ53bWyTYRnAC@J@OUu%bsWyu
z!*|3<t^oQQTgjegOHG%lyug@IW|08iXrihpM+s$09YSxWxR?Mu*Gm5Njpy_7s^wg^
zH`K|DO!;|i*#6v=$sD1dd*f$8Fxb*<LE|~%2XW0OJLXd>v5B|o0r+t%?S9p_#;qwf
z1ucjj=Vfj$?r66RY{h+Z$<*zUCzkmLF$vR?e4g1lF{DK^`*%);+2mV<#d%r#9K%72
zs%m<-j}(`wL~m)iA+=&k=kOw>@HKMTA>h;*Xp@r?JZ`jSgLf<IwH2(**tJEpsMpRf
zC91F)?w8IGeB$xS<BNt+O{fTz9z-j}0;J}r1<-JXRR8o6E9~ugbmiS+K@~WBla&%e
f!$qI9=&$lp6&r<XJDQs#j!ueBL{bB(VhjHQoyMDq

literal 0
HcmV?d00001

diff --git a/tests/unit/utils/test_http.py b/tests/unit/utils/test_http.py
index a2f97701f..2051bcbaf 100644
--- a/tests/unit/utils/test_http.py
+++ b/tests/unit/utils/test_http.py
@@ -5,6 +5,16 @@ 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
@@ -27,5 +37,20 @@ def test_open_url(tmp_path, monkeypatch, http):
 
     with open_url((http / "sample.txt").url) as fd:
         # Test various .read() variants
-        assert fd.read(len(text)) == text
+        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
+        assert read_nbytes(fd, len(text), mode="rb") == text
+        assert fd.peek(len(text)) == text
+        assert fd.read() == text
+        assert fd.read() == b""