Fix typing of some methods in fs.iotools and fs.memoryfs file APIs

Author: althonos

fs/iotools.py

@@ -4,13 +4,15 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import array
 import io
 import typing
 from io import SEEK_SET, SEEK_CUR
 
 from .mode import Mode
 
 if typing.TYPE_CHECKING:
+    import mmap
     from io import RawIOBase
     from typing import (
         Any,
@@ -88,8 +90,11 @@ def truncate(self, size=None):
         return self._f.truncate(size)
 
     def write(self, data):
-        # type: (bytes) -> int
-        count = self._f.write(data)
+        # type: (Union[bytes, bytearray, memoryview, array.array[Any], mmap.mmap]) -> int
+        if isinstance(data, array.array):
+            count = self._f.write(data.tobytes())
+        else:
+            count = self._f.write(data)  # type: ignore
         return len(data) if count is None else count
 
     @typing.no_type_check
@@ -130,17 +135,20 @@ def readinto1(self, b):
             b[:bytes_read] = data
             return bytes_read
 
-    def readline(self, limit=-1):
-        # type: (int) -> bytes
-        return self._f.readline(limit)
+    def readline(self, limit=None):
+        # type: (Optional[int]) -> bytes
+        return self._f.readline(-1 if limit is None else limit)
 
-    def readlines(self, hint=-1):
-        # type: (int) -> List[bytes]
-        return self._f.readlines(hint)
+    def readlines(self, hint=None):
+        # type: (Optional[int]) -> List[bytes]
+        return self._f.readlines(-1 if hint is None else hint)
 
-    def writelines(self, sequence):
-        # type: (Iterable[Union[bytes, bytearray]]) -> None
-        return self._f.writelines(sequence)
+    def writelines(self, lines):
+        # type: (Iterable[Union[bytes, bytearray, memoryview, array.array[Any], mmap.mmap]]) -> None
+        _lines = (
+            line.tobytes() if isinstance(line, array.array) else line for line in lines
+        )
+        return self._f.writelines(typing.cast("Iterable[bytes]", _lines))
 
     def __iter__(self):
         # type: () -> Iterator[bytes]

fs/memoryfs.py

@@ -24,11 +24,14 @@
 from ._typing import overload
 
 if typing.TYPE_CHECKING:
+    import array
+    import mmap
     from typing import (
         Any,
         BinaryIO,
         Collection,
         Dict,
+        Iterable,
         Iterator,
         List,
         Optional,
@@ -116,8 +119,8 @@ def next(self):
 
     __next__ = next
 
-    def readline(self, size=-1):
-        # type: (int) -> bytes
+    def readline(self, size=None):
+        # type: (Optional[int]) -> bytes
         if not self._mode.reading:
             raise IOError("File not open for reading")
         with self._seek_lock():
@@ -131,7 +134,7 @@ def close(self):
                 self._dir_entry.remove_open_file(self)
                 super(_MemoryFile, self).close()
 
-    def read(self, size=-1):
+    def read(self, size=None):
         # type: (Optional[int]) -> bytes
         if not self._mode.reading:
             raise IOError("File not open for reading")
@@ -190,15 +193,15 @@ def writable(self):
         return self._mode.writing
 
     def write(self, data):
-        # type: (bytes) -> int
+        # type: (Union[bytes, bytearray, memoryview, array.array[Any], mmap.mmap]) -> int
         if not self._mode.writing:
             raise IOError("File not open for writing")
         with self._seek_lock():
             self.on_modify()
             return self._bytes_io.write(data)
 
-    def writelines(self, sequence):  # type: ignore
-        # type: (List[bytes]) -> None
+    def writelines(self, sequence):
+        # type: (Iterable[Union[bytes, bytearray, memoryview, array.array[Any], mmap.mmap]]) -> None
         # FIXME(@althonos): For some reason the stub for IOBase.writelines
         #      is List[Any] ?! It should probably be Iterable[ByteString]
         with self._seek_lock():

fs/tree.py

@@ -133,7 +133,7 @@ def format_directory(path, levels):
         try:
             directory = sorted(
                 fs.filterdir(path, exclude_dirs=exclude, files=filter),
-                key=sort_key_dirs_first if dirs_first else sort_key,
+                key=sort_key_dirs_first if dirs_first else sort_key,  # type: ignore
             )
         except Exception as error:
             prefix = (