Binary file handles

neumond · neumond · commit 762021a47fb6 · 2020-07-04T00:36:50.000+03:00
diff --git a/computercraft/subapis/fs.py b/computercraft/subapis/fs.py
@@ -1,9 +1,10 @@
 from contextlib import asynccontextmanager
-from typing import Optional, List
+from typing import Optional, List, Union
 
 from .base import BaseSubAPI
+from ..errors import LuaException
 from ..lua import lua_args
-from ..rproc import boolean, string, integer, nil, array_string, option_string, fact_scheme_dict
+from ..rproc import boolean, string, integer, nil, array_string, option_string, option_integer, fact_scheme_dict
 
 
 attribute = fact_scheme_dict({
@@ -14,9 +15,20 @@
 }, {})
 
 
+class SeekMixin:
+    async def seek(self, whence: str = None, offset: int = None) -> int:
+        # whence: set, cur, end
+        r = await self._send('seek', whence, offset)
+        if isinstance(r, list):
+            assert r[0] is False
+            raise LuaException(r[1])
+        return integer(r)
+
+
 class ReadHandle(BaseSubAPI):
-    async def read(self, count: int) -> Optional[str]:
-        return option_string(await self._send('read', count))
+    async def read(self, count: int = None) -> Optional[Union[str, int]]:
+        r = await self._send('read', count)
+        return option_integer(r) if count is None else option_string(r)
 
     async def readLine(self) -> Optional[str]:
         return option_string(await self._send('readLine'))
@@ -34,6 +46,10 @@ async def __anext__(self):
         return line
 
 
+class BinaryReadHandle(ReadHandle, SeekMixin):
+    pass
+
+
 class WriteHandle(BaseSubAPI):
     async def write(self, text: str):
         return nil(await self._send('write', text))
@@ -45,6 +61,10 @@ async def flush(self):
         return nil(await self._send('flush'))
 
 
+class BinaryWriteHandle(WriteHandle, SeekMixin):
+    pass
+
+
 class FSAPI(BaseSubAPI):
     async def list(self, path: str) -> List[str]:
         return array_string(await self._send('list', path))
@@ -103,7 +123,11 @@ async def open(self, path: str, mode: str):
         )
         fin_tpl = '{e}.close()'
         async with self._cc._create_temp_object(create_expr, fin_tpl) as var:
-            yield (ReadHandle if 'r' in mode else WriteHandle)(self._cc, var)
+            if 'b' in mode:
+                hcls = BinaryReadHandle if 'r' in mode else BinaryWriteHandle
+            else:
+                hcls = ReadHandle if 'r' in mode else WriteHandle
+            yield hcls(self._cc, var)
 
     async def find(self, wildcard: str) -> List[str]:
         return array_string(await self._send('find', wildcard))
diff --git a/testmod.py b/testmod.py
@@ -484,6 +484,24 @@ async def test_fs_api(api):
 
     assert await api.fs.exists('tdir/banana') is True
 
+    async with api.fs.open('tdir/binfile', 'wb') as f:
+        assert await f.write('a' * 9) is None
+        assert await f.seek() == 9
+        assert await f.seek('set', 0) == 0
+        assert await f.write('b' * 3) is None
+        assert await f.seek('cur', -1) == 2
+        assert await f.write('c' * 3) is None
+        assert await f.seek('end') == 9
+        assert await f.write('d' * 3) is None
+        with assert_raises(LuaException):
+            await f.seek('set', -10)
+
+    async with api.fs.open('tdir/binfile', 'rb') as f:
+        assert await f.readAll() == 'bbcccaaaaddd'
+
+    async with api.fs.open('tdir/binfile', 'rb') as f:
+        assert isinstance(await f.read(), int)
+
     assert await api.fs.delete('tdir') is None
     assert await api.fs.delete('tfile') is None
     assert await api.fs.delete('doesnotexist') is None
