From cbf55f199b942581e0cb5ef9d0824394e1d5f47d Mon Sep 17 00:00:00 2001
From: Stanislav <39829332+StannisMod@users.noreply.github.com>
Date: Wed, 15 Nov 2023 12:59:32 +0300
Subject: [PATCH] Fixed merged_stages_datetime_df (#66)

---
 sampo/schemas/schedule.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sampo/schemas/schedule.py b/sampo/schemas/schedule.py
index 6c9baa19..cbcb7d12 100644
--- a/sampo/schemas/schedule.py
+++ b/sampo/schemas/schedule.py
@@ -99,14 +99,14 @@ def _deserialize(cls, dict_representation: T) -> 'Schedule':
         dict_representation['works'] = [ScheduledWork._deserialize(sw) for sw in dict_representation['works']]
         return Schedule.from_scheduled_works(**dict_representation)
 
-    # @lru_cache
+    @lru_cache
     def merged_stages_datetime_df(self, offset: Union[datetime, str]) -> DataFrame:
         """
         Merges split stages of same works after lag optimization and returns schedule DataFrame shifted to start.
         :param offset: Start of schedule, to add as an offset.
         :return: Shifted schedule DataFrame with merged tasks.
         """
-        result = self.offset_schedule(offset)
+        result = fix_split_tasks(self.offset_schedule(offset))
         return result
 
     def offset_schedule(self, offset: Union[datetime, str]) -> DataFrame: