changeset 3087:a51f7fce1e2c

tools (stream): data modification on SatFile: - if `data_cb` is used and if it returns a not None value, it is used instead of the data read from the file. This allows data modification on the fly, useful notably for encryption - new `check_size_with_read` argument which check size on `close()` using amount of data actually read/written instead of file size. This avoid a warning when file is modified on the fly - added `closed` attribute
author Goffi <goffi@goffi.org>
date Fri, 20 Dec 2019 12:28:04 +0100
parents 13be04a70e2f
children d1464548055a
files sat/tools/stream.py
diffstat 1 files changed, 48 insertions(+), 14 deletions(-) [+]
line wrap: on
line diff
--- a/sat/tools/stream.py	Fri Dec 20 12:28:04 2019 +0100
+++ b/sat/tools/stream.py	Fri Dec 20 12:28:04 2019 +0100
@@ -19,14 +19,14 @@
 
 """ interfaces """
 
+import uuid
+import os
+from zope import interface
 from sat.core import exceptions
 from sat.core.constants import Const as C
 from sat.core.log import getLogger
 from twisted.protocols import basic
 from twisted.internet import interfaces
-from zope import interface
-import uuid
-import os
 
 log = getLogger(__name__)
 
@@ -40,28 +40,32 @@
         pass
 
 
-class SatFile(object):
+class SatFile:
     """A file-like object to have high level files manipulation"""
 
     # TODO: manage "with" statement
 
     def __init__(self, host, client, path, mode="rb", uid=None, size=None, data_cb=None,
-                 auto_end_signals=True):
+                 auto_end_signals=True, check_size_with_read=False):
         """
         @param host: %(doc_host)s
-        @param path(str): path of the file to get
+        @param path(Path, str): path to the file to get or write to
         @param mode(str): same as for built-in "open" function
         @param uid(unicode, None): unique id identifing this progressing element
             This uid will be used with self.host.progressGet
             will be automaticaly generated if None
         @param size(None, int): size of the file (when known in advance)
         @param data_cb(None, callable): method to call on each data read/write
-            mainly useful to do things like calculating hash
+            can be used to do processing like calculating hash.
+            if data_cb return a non None value, it will be used instead of the
+            data read/to write
         @param auto_end_signals(bool): if True, progressFinished and progressError signals
             are automatically sent.
             if False, you'll have to call self.progressFinished and self.progressError
             yourself.
             progressStarted signal is always sent automatically
+        @param check_size_with_read(bool): if True, size well be checked using number of
+            bytes read or written. This is useful when data_cb modifiy len of file.
         """
         self.host = host
         self.profile = client.profile
@@ -76,6 +80,19 @@
         )
         self.host.bridge.progressStarted(self.uid, metadata, client.profile)
 
+        self._transfer_count = 0 if check_size_with_read else None
+
+    @property
+    def check_size_with_read(self):
+        return self._transfer_count is not None
+
+    @check_size_with_read.setter
+    def check_size_with_read(self, value):
+        if value and self._transfer_count is None:
+            self._transfer_count = 0
+        else:
+            self._transfer_count = None
+
     def checkSize(self):
         """Check that current size correspond to given size
 
@@ -83,7 +100,10 @@
         @return (bool): True if the position is the same as given size
         @raise exceptions.NotFound: size has not be specified
         """
-        position = self._file.tell()
+        if self.check_size_with_read:
+            position = self._transfer_count
+        else:
+            position = self._file.tell()
         if self.size is None:
             raise exceptions.NotFound
         return position == self.size
@@ -91,8 +111,10 @@
     def close(self, progress_metadata=None, error=None):
         """Close the current file
 
-        @param progress_metadata(None, dict): metadata to send with _onProgressFinished message
-        @param error(None, unicode): set to an error message if progress was not successful
+        @param progress_metadata(None, dict): metadata to send with _onProgressFinished
+            message
+        @param error(None, unicode): set to an error message if progress was not
+            successful
             mutually exclusive with progress_metadata
             error can happen even if error is None, if current size differ from given size
         """
@@ -119,6 +141,10 @@
 
         self.host.removeProgressCb(self.uid, self.profile)
 
+    @property
+    def closed(self):
+        return self._file.closed
+
     def progressFinished(self, metadata=None):
         if metadata is None:
             metadata = {}
@@ -131,14 +157,22 @@
         self._file.flush()
 
     def write(self, buf):
+        if self.data_cb is not None:
+            ret = self.data_cb(buf)
+            if ret is not None:
+                buf = ret
+        if self._transfer_count is not None:
+            self._transfer_count += len(buf)
         self._file.write(buf)
-        if self.data_cb is not None:
-            return self.data_cb(buf)
 
     def read(self, size=-1):
         read = self._file.read(size)
-        if self.data_cb is not None and read:
-            self.data_cb(read)
+        if self.data_cb is not None:
+            ret = self.data_cb(read)
+            if ret is not None:
+                read = ret
+        if self._transfer_count is not None:
+            self._transfer_count += len(read)
         return read
 
     def seek(self, offset, whence=os.SEEK_SET):