From 9ddfd9ab1e0744f27cbe132d6d8d967dfb0767be Mon Sep 17 00:00:00 2001
From: Fredrik Tolf <fredrik@dolda2000.com>
Date: Sun, 11 Aug 2024 04:17:31 +0200
Subject: [PATCH] Make limitreader more compatible with general io classes.

---
 wrw/req.py | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/wrw/req.py b/wrw/req.py
index 010e907..d5e22a3 100644
--- a/wrw/req.py
+++ b/wrw/req.py
@@ -56,9 +56,10 @@ class shortinput(IOError, EOFError):
         super().__init__("Unexpected EOF")
 
 class limitreader(object):
-    def __init__(self, back, limit):
+    def __init__(self, back, limit, short=False):
         self.bk = back
         self.limit = limit
+        self.short = short
         self.rb = 0
         self.buf = bytearray()
 
@@ -72,11 +73,15 @@ class limitreader(object):
         while len(self.buf) < ra:
             ret = self.bk.read(ra - len(self.buf))
             if ret == b"":
+                if self.short:
+                    ret = bytes(self.buf)
+                    self.buf[:] = b""
+                    return ret
                 raise shortinput()
             self.buf.extend(ret)
             self.rb += len(ret)
         ret = bytes(self.buf[:ra])
-        self.buf = self.buf[ra:]
+        self.buf[:ra] = b""
         return ret
 
     def readline(self, size=-1):
@@ -85,16 +90,16 @@ class limitreader(object):
             p = self.buf.find(b'\n', off)
             if p >= 0:
                 ret = bytes(self.buf[:p + 1])
-                self.buf = self.buf[p + 1:]
+                self.buf[:p + 1] = b""
                 return ret
             off = len(self.buf)
             if size >= 0 and len(self.buf) >= size:
                 ret = bytes(self.buf[:size])
-                self.buf = self.buf[size:]
+                self.buf[:size] = b""
                 return ret
             if self.rb == self.limit:
                 ret = bytes(self.buf)
-                self.buf = bytearray()
+                self.buf[:] = b""
                 return ret
             ra = self.limit - self.rb
             if size >= 0:
@@ -102,6 +107,10 @@ class limitreader(object):
             ra = min(ra, 1024)
             ret = self.bk.read(ra)
             if ret == b"":
+                if self.short:
+                    ret = bytes(self.buf)
+                    self.buf[:] = b""
+                    return ret
                 raise shortinput()
             self.buf.extend(ret)
             self.rb += len(ret)
@@ -120,6 +129,16 @@ class limitreader(object):
                 return ret
         return lineiter()
 
+    def readable(self):
+        return True
+    def writable(self):
+        return False
+    def seekable(self):
+        return False
+    @property
+    def closed(self):
+        return self.bk.closed
+
 class request(object):
     def copy(self):
         return copyrequest(self)
-- 
2.11.0