From 5d39784abb1228bddbb5abf0fc329c5d3d4aecfc Mon Sep 17 00:00:00 2001
From: Jakub Jezek <jakub@orbi.tools>
Date: Wed, 5 Jan 2022 22:32:07 +0100
Subject: [PATCH] flame: maintained segment selection

---
 openpype/hosts/flame/api/lib.py      | 45 +++++++++++++++++++++++++---
 openpype/hosts/flame/api/pipeline.py | 26 ----------------
 2 files changed, 41 insertions(+), 30 deletions(-)

diff --git a/openpype/hosts/flame/api/lib.py b/openpype/hosts/flame/api/lib.py
index f91f593eb5..5860bb728d 100644
--- a/openpype/hosts/flame/api/lib.py
+++ b/openpype/hosts/flame/api/lib.py
@@ -393,15 +393,13 @@ def set_segment_pype_tag(segment, data=None):
         tag_data.update(data)
         # update marker with tag data
         marker.comment = json.dumps(tag_data)
-
-        return True
     else:
         # update tag data with new data
         marker = create_pype_marker(segment)
         # add tag data to marker's comment
         marker.comment = json.dumps(data)
 
-        return True
+    return True
 
 
 
@@ -486,4 +484,43 @@ def create_pype_marker(segment):
     # set colour
     marker.colour = ctx.marker_color
 
-    return marker
\ No newline at end of file
+    return marker
+
+
+@contextlib.contextmanager
+def maintained_segment_selection(sequence):
+    """Maintain selection during context
+
+    Example:
+        >>> with maintained_selection():
+        ...     node['selected'].setValue(True)
+        >>> print(node['selected'].value())
+        False
+    """
+    selected_segments = []
+    for ver in sequence.versions:
+        for track in ver.tracks:
+            if len(track.segments) == 0 and track.hidden:
+                continue
+            for segment in track.segments:
+                if segment.selected != True:
+                    continue
+                selected_segments.append(segment)
+    try:
+        # do the operation
+        yield
+    finally:
+        reset_segment_selection(sequence)
+        for segment in selected_segments:
+            segment.selected = True
+
+
+def reset_segment_selection(sequence):
+    """Deselect all selected nodes
+    """
+    for ver in sequence.versions:
+        for track in ver.tracks:
+            if len(track.segments) == 0 and track.hidden:
+                continue
+            for segment in track.segments:
+                segment.selected = False
diff --git a/openpype/hosts/flame/api/pipeline.py b/openpype/hosts/flame/api/pipeline.py
index 26dfe7c032..00860857f1 100644
--- a/openpype/hosts/flame/api/pipeline.py
+++ b/openpype/hosts/flame/api/pipeline.py
@@ -97,32 +97,6 @@ def update_container(tl_segment, data=None):
     # TODO: update_container
     pass
 
-
-@contextlib.contextmanager
-def maintained_selection():
-    """Maintain selection during context
-
-    Example:
-        >>> with maintained_selection():
-        ...     node['selected'].setValue(True)
-        >>> print(node['selected'].value())
-        False
-    """
-    # TODO: maintained_selection + remove undo steps
-
-    try:
-        # do the operation
-        yield
-    finally:
-        pass
-
-
-def reset_selection():
-    """Deselect all selected nodes
-    """
-    pass
-
-
 def on_pyblish_instance_toggled(instance, old_value, new_value):
     """Toggle node passthrough states on instance toggles."""