Various lock stitch options (#2006)

Co-authored-by: Lex Neva

1 files changed, 17 insertions, 4 deletions

diff --git a/lib/stitch_plan/stitch_group.py b/lib/stitch_plan/stitch_group.py
index 717bb617..957da3f8 100644
--- a/lib/stitch_plan/stitch_group.py
+++ b/lib/stitch_plan/stitch_group.py
@@ -18,7 +18,7 @@ class StitchGroup:
     """
 
     def __init__(self, color=None, stitches=None, trim_after=False, stop_after=False,
-                 tie_modus=0, force_lock_stitches=False, stitch_as_is=False, tags=None):
+                 lock_stitches=(None, None), force_lock_stitches=False, tags=None):
         # DANGER: if you add new attributes, you MUST also set their default
         # values in __new__() below.  Otherwise, cached stitch plans can be
         # loaded and create objects without those properties defined, because
@@ -27,9 +27,8 @@ class StitchGroup:
         self.color = color
         self.trim_after = trim_after
         self.stop_after = stop_after
-        self.tie_modus = tie_modus
+        self.lock_stitches = lock_stitches
         self.force_lock_stitches = force_lock_stitches
-        self.stitch_as_is = stitch_as_is
         self.stitches = []
 
         if stitches:
@@ -44,11 +43,14 @@ class StitchGroup:
         # Set default values for any new attributes here (see note in __init__() above)
         # instance.foo = None
 
+        instance.lock_stitches = None
+        
         return instance
 
     def __add__(self, other):
         if isinstance(other, StitchGroup):
-            return StitchGroup(self.color, self.stitches + other.stitches)
+            return StitchGroup(self.color, self.stitches + other.stitches,
+                               lock_stitches=self.lock_stitches, force_lock_stitches=self.force_lock_stitches)
         else:
             raise TypeError("StitchGroup can only be added to another StitchGroup")
 
@@ -77,3 +79,14 @@ class StitchGroup:
     def add_tag(self, tag):
         for stitch in self.stitches:
             stitch.add_tag(tag)
+
+    def get_lock_stitches(self, pos, disable_ties=False):
+        if len(self.stitches) < 2:
+            return []
+
+        lock_pos = 0 if pos == "start" else 1
+        if disable_ties or self.lock_stitches[lock_pos] is None:
+            return
+
+        stitches = self.lock_stitches[lock_pos].stitches(self.stitches, pos)
+        return stitches