From 99c89f38118f9d870a5df8549604a3cf97ff5416 Mon Sep 17 00:00:00 2001
From: Milan Kolar <milan@orbi.tools>
Date: Tue, 3 Dec 2019 18:43:11 +0100
Subject: [PATCH] add option for start and end frame to burnins

---
 pype/plugins/global/publish/extract_burnin.py | 15 +++++--
 pype/scripts/otio_burnin.py                   | 40 +++++++++++++++----
 2 files changed, 44 insertions(+), 11 deletions(-)

diff --git a/pype/plugins/global/publish/extract_burnin.py b/pype/plugins/global/publish/extract_burnin.py
index b1569aaa45..457f4801cf 100644
--- a/pype/plugins/global/publish/extract_burnin.py
+++ b/pype/plugins/global/publish/extract_burnin.py
@@ -29,11 +29,16 @@ class ExtractBurnin(pype.api.Extractor):
         if instance.context.data.get('version'):
             version = "v" + str(instance.context.data['version'])
 
+        frame_start = int(instance.data.get("frameStart") or 0)
+        frame_end = int(instance.data.get("frameEnd") or 1)
+        duration = frame_end - frame_start + 1
         prep_data = {
             "username": instance.context.data['user'],
             "asset": os.environ['AVALON_ASSET'],
             "task": os.environ['AVALON_TASK'],
-            "start_frame": int(instance.data["frameStart"]),
+            "frame_start": frame_start,
+            "frame_end": frame_end,
+            "duration": duration,
             "version": version
         }
         self.log.debug("__ prep_data: {}".format(prep_data))
@@ -49,8 +54,12 @@ class ExtractBurnin(pype.api.Extractor):
             name = "_burnin"
             movieFileBurnin = filename.replace(".mov", "") + name + ".mov"
 
-            full_movie_path = os.path.join(os.path.normpath(stagingdir), repre["files"])
-            full_burnin_path = os.path.join(os.path.normpath(stagingdir), movieFileBurnin)
+            full_movie_path = os.path.join(
+                os.path.normpath(stagingdir), repre["files"]
+            )
+            full_burnin_path = os.path.join(
+                os.path.normpath(stagingdir), movieFileBurnin
+            )
             self.log.debug("__ full_burnin_path: {}".format(full_burnin_path))
 
             burnin_data = {
diff --git a/pype/scripts/otio_burnin.py b/pype/scripts/otio_burnin.py
index ad2e59fc96..a7ae73f6c3 100644
--- a/pype/scripts/otio_burnin.py
+++ b/pype/scripts/otio_burnin.py
@@ -91,7 +91,9 @@ class ModifiedBurnins(ffmpeg_burnins.Burnins):
         text = today.strftime(date_format)
         self._add_burnin(text, align, options, ffmpeg_burnins.DRAWTEXT)
 
-    def add_frame_numbers(self, align, options=None, start_frame=None):
+    def add_frame_numbers(
+        self, align, options=None, start_frame=None, text=None
+    ):
         """
         Convenience method to create the frame number expression.
 
@@ -103,7 +105,12 @@ class ModifiedBurnins(ffmpeg_burnins.Burnins):
         if start_frame:
             options['frame_offset'] = start_frame
 
-        options['expression'] = r'%%{eif\:n+%d\:d}' % options['frame_offset']
+        expr = r'%%{eif\:n+%d\:d}' % options['frame_offset']
+        if text and isinstance(text, str):
+            text = r"{}".format(text)
+            expr = text.replace("{current_frame}", expr)
+
+        options['expression'] = expr
         text = str(int(self.end_frame + options['frame_offset']))
         self._add_burnin(text, align, options, ffmpeg_burnins.DRAWTEXT)
 
@@ -121,7 +128,7 @@ class ModifiedBurnins(ffmpeg_burnins.Burnins):
 
         timecode = ffmpeg_burnins._frames_to_timecode(
             options['frame_offset'],
-           self.frame_rate
+            self.frame_rate
         )
         options = options.copy()
         if not options.get('fps'):
@@ -284,8 +291,8 @@ def burnins_from_data(input_path, output_path, data, overwrite=True):
 
     burnin = ModifiedBurnins(input_path, options_init=options_init)
 
-    start_frame = data.get("start_frame")
-    start_frame_tc = data.get('start_frame_tc', start_frame)
+    frame_start = data.get("frame_start")
+    frame_start_tc = data.get('frame_start_tc', frame_start)
     for align_text, preset in presets.get('burnins', {}).items():
         align = None
         if align_text == 'TOP_LEFT':
@@ -311,7 +318,7 @@ def burnins_from_data(input_path, output_path, data, overwrite=True):
 
         if (
             bi_func in ['frame_numbers', 'timecode'] and
-            start_frame is None
+            frame_start is None
         ):
             log.error(
                 'start_frame is not set in entered data!'
@@ -320,9 +327,26 @@ def burnins_from_data(input_path, output_path, data, overwrite=True):
             return
 
         if bi_func == 'frame_numbers':
-            burnin.add_frame_numbers(align, start_frame=start_frame)
+            current_frame_identifier = "{current_frame}"
+            text = preset.get('text') or current_frame_identifier
+
+            if current_frame_identifier not in text:
+                log.warning((
+                    'Text for Frame numbers don\'t have '
+                    '`{current_frame}` key in text!'
+                ))
+
+            text_items = []
+            split_items = text.split(current_frame_identifier)
+            for item in split_items:
+                text_items.append(item.format(**data))
+
+            text = "{current_frame}".join(text_items)
+
+            burnin.add_frame_numbers(align, start_frame=frame_start, text=text)
+
         elif bi_func == 'timecode':
-            burnin.add_timecode(align, start_frame=start_frame_tc)
+            burnin.add_timecode(align, start_frame=frame_start_tc)
         elif bi_func == 'text':
             if not preset.get('text'):
                 log.error('Text is not set for text function burnin!')