Fixed merged_stages_datetime_df (#66)

aimclub · Nov 15, 2023 · cbf55f1 · cbf55f1
1 parent 8a8ab28
commit cbf55f1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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: