Also handle size argument in readlines

sebastinas · sebastinas · commit 6bdeadde6b67 · 2022-05-28T17:21:14.000+02:00
diff --git a/bpython/curtsiesfrontend/repl.py b/bpython/curtsiesfrontend/repl.py
@@ -193,8 +193,22 @@ def readline(self, size: int = -1) -> str:
         self.readline_results.append(value)
         return value if size <= -1 else value[:size]
 
-    def readlines(self, size=-1):
-        return list(iter(self.readline, ""))
+    def readlines(self, size: Optional[int] = -1) -> List[str]:
+        if size is None:
+            # the default readlines implementation also accepts None
+            size = -1
+        if not isinstance(size, int):
+            raise TypeError("argument should be integer or None, not 'str'")
+        if size <= 0:
+            # read as much as we can
+            return list(iter(self.readline, ""))
+
+        lines = []
+        while size > 0:
+            line = self.readline()
+            lines.append(line)
+            size -= len(line)
+        return lines
 
     def __iter__(self):
         return iter(self.readlines())
