widgets/thread: fix processing quote blocks into folding tree

1 files changed, 11 insertions, 6 deletions

diff --git a/alot/widgets/thread.py b/alot/widgets/thread.py
index a1512679..b0276eb0 100644
--- a/alot/widgets/thread.py
+++ b/alot/widgets/thread.py
@@ -349,7 +349,7 @@ class _TextPart(_MIMEPartWidget):
                     b['level'] = merge_level
 
         fold_stack = []
-        cur_fold   = _Fold(0, 0, None)
+        cur_fold   = _Fold(0, 0, 0)
         max_level  = 0
         for b in blocks:
             max_level = max(max_level, b['level'])
@@ -357,14 +357,19 @@ class _TextPart(_MIMEPartWidget):
             if b['level'] > cur_fold.level:
                 fold_stack.append(cur_fold)
                 cur_fold = _Fold(b['level'], b['start'], b['end'])
-            elif b['level'] == cur_fold.level:
-                cur_fold.end = b['end']
             else:
-                child    = cur_fold
+                while b['level'] < cur_fold.level:
+                    child    = cur_fold
+
+                    if fold_stack[-1].level < cur_fold.level - 1:
+                        cur_fold = _Fold(cur_fold.level - 1, cur_fold.start, cur_fold.end)
+                    else:
+                        cur_fold = fold_stack.pop()
+
+                    cur_fold.children.append(child)
+                    cur_fold.end = max(cur_fold.end, child.end)
 
-                cur_fold     = fold_stack.pop()
                 cur_fold.end = b['end']
-                cur_fold.children.append(child)
 
         while fold_stack:
             fold_stack[-1].children.append(cur_fold)