scripts: ftp: Implement file-like object for IO.

Issue #128.

mavros/scripts/mavftp

@@ -30,11 +30,6 @@ from mavros.srv import FileOpen, FileClose, FileRead, FileList, FileOpenRequest,
     FileMakeDir, FileRemoveDir, FileRemove, FileWrite
 
 
-def print_if(cond, *args, **kvargs):
-    if cond:
-        print(*args, **kvargs)
-
-
 def fault(*args, **kvargs):
     kvargs['file'] = sys.stdout
     print(*args, **kvargs)
@@ -46,6 +41,117 @@ def check_ret(operation, ret):
         fault(operation, "failed:", os.strerror(ret.r_errno))
 
 
+class FTPFile(object):
+    def __init__(self, name, mode, ns="/mavros"):
+        self.name = None
+        self.mode = mode
+        self.mavros_ns = ns
+        self.open(name, mode)
+
+    def __del__(self):
+        self.close()
+
+    def _check_raise_errno(self, ret):
+        if not ret.success:
+            raise IOError(os.strerror(ret.r_errno))
+
+    def open(self, path, mode):
+        """
+        Supported modes:
+            - 'w': write binary
+            - 'r': read binary
+            - 'cw': create excl & write
+        """
+        if mode == 'w' or mode == 'wb':
+            m = FileOpenRequest.MODE_WRITE
+        elif mode == 'r' or mode == 'rb':
+            m = FileOpenRequest.MODE_READ
+        elif mode == 'cw':
+            m = FileOpenRequest.MODE_CREATE
+        else:
+            fault("Unknown open mode: ", m)
+
+        try:
+            open_cl = rospy.ServiceProxy(self.mavros_ns + "/ftp/open", FileOpen)
+            ret = open_cl(file_path=path, mode=m)
+        except rospy.ServiceException as ex:
+            raise IOError(str(ex))
+
+        self._check_raise_errno(ret)
+
+        self._read_cl = rospy.ServiceProxy(self.mavros_ns + "/ftp/read", FileRead)
+        self._write_cl = rospy.ServiceProxy(self.mavros_ns + "/ftp/write", FileWrite)
+
+        self.name = path
+        self.mode = mode
+        self.size = ret.size
+        self.offset = 0
+
+    def close(self):
+        if self.closed:
+            return
+
+        try:
+            close_cl = rospy.ServiceProxy(self.mavros_ns + "/ftp/close", FileClose)
+            ret = close_cl(file_path=self.name)
+        except rospy.ServiceException as ex:
+            raise IOError(str(ex))
+
+        self.name = None
+        self._check_raise_errno(ret)
+
+    def read(self, size=1):
+        try:
+            ret = self._read_cl(file_path=self.name, offset=self.offset, size=size)
+        except rospy.ServiceException as ex:
+            raise IOError(str(ex))
+
+        self._check_raise_errno(ret)
+        self.offset += len(ret.data)
+        return bytearray(ret.data)
+
+    def write(self, bin_data):
+        data_len = len(bin_data)
+        try:
+            ret = self._write_cl(file_path=self.name, offset=self.offset, data=bin_data)
+        except rospy.ServiceException as ex:
+            raise IOError(str(ex))
+
+        self._check_raise_errno(ret)
+        self.offset += data_len
+        if self.offset > self.size:
+            self.size = self.offset
+
+    def tell(self):
+        return self.offset
+
+    def seek(self, offset, whence=os.SEEK_SET):
+        if whence is os.SEEK_SET:
+            self.offset = offset
+        elif whence is os.SEEK_END:
+            self.offset = offset + self.size
+        elif whence is os.SEEK_CUR:
+            self.offset += offset
+        else:
+            raise ValueError("Unknown whence")
+
+    def truncate(self, size=0):
+        raise NotImplementedError
+
+    @property
+    def closed(self):
+        return self.name is None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
+
+# -*- subcommand handlers -*-
+
+
 def do_list(args):
     rospy.init_node("mavftp", anonymous=True)
 
@@ -67,54 +173,11 @@ def do_list(args):
         print(n)
 
 
-def open_file(args, path, mode):
-    if mode == 'w':
-        m = FileOpenRequest.MODE_WRITE
-    else:
-        m = FileOpenRequest.MODE_READ
-
-    try:
-        open_cl = rospy.ServiceProxy(args.mavros_ns + "/ftp/open", FileOpen)
-        ret = open_cl(file_path=path, mode=m)
-    except rospy.ServiceException as ex:
-        fault(ex)
-
-    return (ret, ret.size)
-
-
-def close_file(args, path):
-    try:
-        close_cl = rospy.ServiceProxy(args.mavros_ns + "/ftp/close", FileClose)
-        ret = close_cl(file_path=path)
-    except rospy.ServiceException as ex:
-        fault(ex)
-
-    return ret
-
-
-def read_file(args, path, offset, size):
-    try:
-        read_cl = rospy.ServiceProxy(args.mavros_ns + "/ftp/read", FileRead)
-        ret = read_cl(file_path=path, offset=offset, size=size)
-    except rospy.ServiceException as ex:
-        fault(ex)
-
-    return (ret, bytearray(ret.data))
-
-
 def do_cat(args):
     rospy.init_node("mavftp", anonymous=True)
 
-    ret, size = open_file(args, args.path, 'r')
-    check_ret("Open", ret)
-
-    try:
-        ret, data = read_file(args, args.path, 0, size)
-        check_ret("Read", ret)
-    finally:
-        close_file(args, args.path)
-
-    if ret.success:
+    with FTPFile(args.path, 'r', args.mavros_ns) as fd:
+        data = fd.read(fd.size)
         sys.stdout.write(data)