[path-] fix undercounted progress for multibyte chars

visidata/path.py

@@ -95,14 +95,15 @@ def peek(self, size=-1):
 
 
 class FileProgress:
-    'Open file in binary mode and track read() progress.'
+    'Open file as text or binary, then track read() progress.'
     def __init__(self, path, fp, mode='r', **kwargs):
         self.path = path
         self.fp = fp
         self.prog = None
         if 'r' in mode:
             gerund = 'reading'
-            self.prog = Progress(gerund=gerund, total=filesize(path))
+            self.total = filesize(path)
+            self.prog = Progress(gerund=gerund, total=self.total)
         elif 'w' in mode:
             gerund = 'writing'
             self.prog = Progress(gerund=gerund)
@@ -117,6 +118,14 @@ def __init__(self, path, fp, mode='r', **kwargs):
         self.fp.read = self.read
         self.fp.close = self.close
 
+        self.est_charbytes = 0
+        self.sampled = ''
+        self.est_ctr = 0
+        self.est_interval = 1000
+        self.est_total = 0
+
+        self.readline_chunks = []
+
         if self.prog:
             self.prog.__enter__()
 
@@ -126,17 +135,46 @@ def close(self, *args, **kwargs):
             self.prog = None
         return self.fp_orig_close(*args, **kwargs)
 
+    def update_progress(self, r):
+        if not self.prog: return
+        if isinstance(r, bytes):
+            r_bytes = len(r)
+            self.prog.addProgress(r_bytes)
+        else:
+            self.update_estimate(r)
+            r_bytes = len(r) * self.est_charbytes
+            self.est_total += r_bytes
+            if self.est_total <= self.total:
+                self.prog.addProgress(r_bytes)
+
+    def update_progress_line(self, line):
+        # updating progress for each line is slow, so do it in batches
+        self.readline_chunks.append(line)
+        if len(self.readline_chunks) == 1000:
+            joiner = b'\n' if isinstance(line, bytes) else '\n'
+            batched = joiner.join(self.readline_chunks)
+            self.update_progress(batched)
+            self.readline_chunks = []
+
+    def update_estimate(self, r):
+        # A short string can cause charbytes to be overestimated by 30%,
+        # due to the Byte Order Marker in encodings like utf-8-sig.
+        # Combining short strings into one big one lowers that error to < 1%.
+        if self.est_ctr % self.est_interval == 0 or len(self.sampled) < 1000:
+            sample = r[:100]
+            self.sampled += sample
+            self.est_charbytes = len(self.sampled.encode(self.encoding)) / len(self.sampled)
+        self.est_ctr += 1
+
     def read(self, size=-1):
         r = self.fp_orig_read(size)
-        if self.prog:
-            if r:
-                self.prog.addProgress(len(r))
+        if r:
+            self.update_progress(r)
         return r
 
     def readline(self, size=-1):
         r = self.fp_orig_readline(size)
-        if self.prog:
-            self.prog.addProgress(len(r))
+        self.update_progress_line(r)
         return r
 
     def __getattr__(self, k):
@@ -148,15 +186,15 @@ def __enter__(self):
 
     def __next__(self):
         r = next(self.fp)
-        self.prog.addProgress(len(r))
+        self.update_progress(r)
         return r
 
     def __iter__(self):
         if not self.prog:
             yield from self.fp
         else:
             for line in self.fp:
-                self.prog.addProgress(len(line))
+                self.update_progress_line(line)
                 yield line
 
     def __exit__(self, type, value, tb):