From 208a9345ae1896c7df8c5c7e44941cbc14ed0a26 Mon Sep 17 00:00:00 2001 From: Tamas Nepusz Date: Sat, 21 Oct 2023 12:46:44 +0200 Subject: [PATCH] refactor: added get_storyboard() function --- src/modules/sbstudio/plugin/model/storyboard.py | 9 +++++++++ .../plugin/operators/add_markers_from_zipped_csv.py | 5 +++-- src/modules/sbstudio/plugin/operators/base.py | 11 +++++------ .../plugin/operators/create_takeoff_grid.py | 3 ++- src/modules/sbstudio/plugin/operators/land.py | 10 ++++++---- .../plugin/operators/move_storyboard_entry.py | 6 ++++-- .../sbstudio/plugin/operators/remove_formation.py | 3 ++- .../operators/remove_schedule_override_entry.py | 6 +++++- .../plugin/operators/remove_storyboard_entry.py | 3 ++- .../sbstudio/plugin/operators/return_to_home.py | 13 ++++++------- .../plugin/operators/select_storyboard_entry.py | 7 +++---- src/modules/sbstudio/plugin/operators/takeoff.py | 6 ++++-- .../sbstudio/plugin/panels/transition_editor.py | 11 ++++------- src/modules/sbstudio/plugin/props/frame_range.py | 9 ++++----- 14 files changed, 59 insertions(+), 43 deletions(-) diff --git a/src/modules/sbstudio/plugin/model/storyboard.py b/src/modules/sbstudio/plugin/model/storyboard.py index c4e30817..29c13af7 100644 --- a/src/modules/sbstudio/plugin/model/storyboard.py +++ b/src/modules/sbstudio/plugin/model/storyboard.py @@ -786,3 +786,12 @@ def _sort_entries(self) -> None: """Sort the items in the storyboard in ascending order of start time.""" # Sort the items in the storyboard itself sort_collection(self.entries, key=StoryboardEntry.sort_key) + + +@with_context +def get_storyboard(context: Optional[Context] = None) -> Storyboard: + """Helper function to retrieve the storyboard of the add-on from the + given context object. + """ + assert context is not None + return context.scene.skybrush.storyboard diff --git a/src/modules/sbstudio/plugin/operators/add_markers_from_zipped_csv.py b/src/modules/sbstudio/plugin/operators/add_markers_from_zipped_csv.py index bc9b92a0..8c0fb89a 100644 --- a/src/modules/sbstudio/plugin/operators/add_markers_from_zipped_csv.py +++ b/src/modules/sbstudio/plugin/operators/add_markers_from_zipped_csv.py @@ -20,6 +20,7 @@ find_f_curve_for_data_path_and_index, ) from sbstudio.plugin.model.formation import add_points_to_formation +from sbstudio.plugin.model.storyboard import get_storyboard from .base import FormationOperator @@ -73,8 +74,8 @@ def execute_on_formation(self, formation, context): fps = context.scene.render.fps # try to figure out the start frame of this formation - storyboard_entry = ( - context.scene.skybrush.storyboard.get_first_entry_for_formation(formation) + storyboard_entry = get_storyboard(context).get_first_entry_for_formation( + formation ) frame_start = ( storyboard_entry.frame_start diff --git a/src/modules/sbstudio/plugin/operators/base.py b/src/modules/sbstudio/plugin/operators/base.py index ccb509f5..5fded64a 100644 --- a/src/modules/sbstudio/plugin/operators/base.py +++ b/src/modules/sbstudio/plugin/operators/base.py @@ -17,6 +17,7 @@ add_points_to_formation, get_markers_from_formation, ) +from sbstudio.plugin.model.storyboard import get_storyboard from sbstudio.plugin.props.frame_range import FrameRangeProperty from sbstudio.plugin.selection import select_only from sbstudio.plugin.selection import Collections @@ -78,7 +79,7 @@ def poll(cls, context): return context.scene.skybrush and context.scene.skybrush.storyboard def execute(self, context): - storyboard = context.scene.skybrush.storyboard + storyboard = get_storyboard(context) validate = getattr(self.__class__, "only_with_valid_storyboard", False) @@ -108,7 +109,7 @@ def poll(cls, context): ) def execute(self, context): - entry = context.scene.skybrush.storyboard.active_entry + entry = get_storyboard(context).active_entry return self.execute_on_storyboard_entry(entry, context) @@ -240,10 +241,8 @@ def execute_on_formation(self, formation, context): # Add a light effect containing the colors of the markers if needed if should_import_colors: # try to figure out the start frame of this formation - storyboard_entry = ( - context.scene.skybrush.storyboard.get_first_entry_for_formation( - formation - ) + storyboard_entry = get_storyboard(context).get_first_entry_for_formation( + formation ) frame_start = ( storyboard_entry.frame_start diff --git a/src/modules/sbstudio/plugin/operators/create_takeoff_grid.py b/src/modules/sbstudio/plugin/operators/create_takeoff_grid.py index 47a71c6f..4c9736e9 100644 --- a/src/modules/sbstudio/plugin/operators/create_takeoff_grid.py +++ b/src/modules/sbstudio/plugin/operators/create_takeoff_grid.py @@ -14,6 +14,7 @@ create_keyframe_for_diffuse_color_of_material, ) from sbstudio.plugin.model.formation import add_points_to_formation, create_formation +from sbstudio.plugin.model.storyboard import get_storyboard from sbstudio.plugin.operators.detach_materials_from_template import ( detach_material_from_drone_template, ) @@ -335,7 +336,7 @@ def _run(self, context): # points if there is one takeoff_grid = Formations.find_takeoff_grid(create=False) if not takeoff_grid: - storyboard = context.scene.skybrush.storyboard + storyboard = get_storyboard(context) entry = storyboard.add_new_entry( formation=create_formation("Takeoff grid", points), frame_start=context.scene.frame_start, diff --git a/src/modules/sbstudio/plugin/operators/land.py b/src/modules/sbstudio/plugin/operators/land.py index 17d1f2b1..c47c54ee 100644 --- a/src/modules/sbstudio/plugin/operators/land.py +++ b/src/modules/sbstudio/plugin/operators/land.py @@ -10,6 +10,7 @@ from sbstudio.plugin.constants import Collections from sbstudio.plugin.model.formation import create_formation from sbstudio.plugin.model.safety_check import get_proximity_warning_threshold +from sbstudio.plugin.model.storyboard import get_storyboard from sbstudio.plugin.utils.evaluator import create_position_evaluator from .base import StoryboardOperator @@ -73,8 +74,9 @@ def poll(cls, context): return drones is not None and len(drones.objects) > 0 def invoke(self, context, event): - storyboard = context.scene.skybrush.storyboard - self.start_frame = max(context.scene.frame_current, storyboard.frame_end) + self.start_frame = max( + context.scene.frame_current, get_storyboard(context).frame_end + ) return context.window_manager.invoke_props_dialog(self) def execute_on_storyboard(self, storyboard, entries, context): @@ -178,9 +180,9 @@ def _run(self, storyboard, *, context) -> bool: def _validate_start_frame(self, context: Context) -> bool: """Returns whether the landing time chosen by the user is valid.""" - storyboard = context.scene.skybrush.storyboard + storyboard = get_storyboard(context) if storyboard.last_entry is not None: - last_frame = context.scene.skybrush.storyboard.frame_end + last_frame = storyboard.frame_end else: last_frame = None diff --git a/src/modules/sbstudio/plugin/operators/move_storyboard_entry.py b/src/modules/sbstudio/plugin/operators/move_storyboard_entry.py index 7116d802..0793eb12 100644 --- a/src/modules/sbstudio/plugin/operators/move_storyboard_entry.py +++ b/src/modules/sbstudio/plugin/operators/move_storyboard_entry.py @@ -1,3 +1,5 @@ +from sbstudio.plugin.model.storyboard import get_storyboard + from .base import StoryboardOperator __all__ = ("MoveStoryboardEntryDownOperator", "MoveStoryboardEntryUpOperator") @@ -17,7 +19,7 @@ def poll(cls, context): if not StoryboardOperator.poll(context): return False - storyboard = context.scene.skybrush.storyboard + storyboard = get_storyboard(context) if storyboard.active_entry is None: return False @@ -45,7 +47,7 @@ def poll(cls, context): if not StoryboardOperator.poll(context): return False - storyboard = context.scene.skybrush.storyboard + storyboard = get_storyboard(context) if storyboard.active_entry is None: return False diff --git a/src/modules/sbstudio/plugin/operators/remove_formation.py b/src/modules/sbstudio/plugin/operators/remove_formation.py index 45104141..b13458a5 100644 --- a/src/modules/sbstudio/plugin/operators/remove_formation.py +++ b/src/modules/sbstudio/plugin/operators/remove_formation.py @@ -1,3 +1,4 @@ +from sbstudio.plugin.model.storyboard import get_storyboard from sbstudio.plugin.objects import remove_objects from .base import FormationOperator @@ -14,7 +15,7 @@ class RemoveFormationOperator(FormationOperator): bl_description = "Remove the selected formation from the show" def execute_on_formation(self, formation, context): - storyboard = context.scene.skybrush.storyboard + storyboard = get_storyboard(context) for entry in storyboard.entries: if entry.formation is formation: diff --git a/src/modules/sbstudio/plugin/operators/remove_schedule_override_entry.py b/src/modules/sbstudio/plugin/operators/remove_schedule_override_entry.py index 3a55697b..f4bf6512 100644 --- a/src/modules/sbstudio/plugin/operators/remove_schedule_override_entry.py +++ b/src/modules/sbstudio/plugin/operators/remove_schedule_override_entry.py @@ -1,3 +1,5 @@ +from sbstudio.plugin.model.storyboard import get_storyboard + from .base import StoryboardEntryOperator __all__ = ("RemoveScheduleOverrideEntryOperator",) @@ -19,7 +21,9 @@ def poll(cls, context): if not StoryboardEntryOperator.poll(context): return False - entry = context.scene.skybrush.storyboard.active_entry + entry = get_storyboard(context).active_entry + assert entry is not None + return entry.active_schedule_override_entry is not None def execute_on_storyboard_entry(self, entry, context): diff --git a/src/modules/sbstudio/plugin/operators/remove_storyboard_entry.py b/src/modules/sbstudio/plugin/operators/remove_storyboard_entry.py index b7b256f1..afd5ea84 100644 --- a/src/modules/sbstudio/plugin/operators/remove_storyboard_entry.py +++ b/src/modules/sbstudio/plugin/operators/remove_storyboard_entry.py @@ -1,4 +1,5 @@ from sbstudio.plugin.constants import Collections +from sbstudio.plugin.model.storyboard import get_storyboard from sbstudio.plugin.transition import find_transition_constraint_between from .base import StoryboardOperator @@ -17,7 +18,7 @@ class RemoveStoryboardEntryOperator(StoryboardOperator): def poll(cls, context): return ( StoryboardOperator.poll(context) - and context.scene.skybrush.storyboard.active_entry is not None + and get_storyboard(context).active_entry is not None ) def execute_on_storyboard(self, storyboard, context): diff --git a/src/modules/sbstudio/plugin/operators/return_to_home.py b/src/modules/sbstudio/plugin/operators/return_to_home.py index 9b53e9a1..8465b949 100644 --- a/src/modules/sbstudio/plugin/operators/return_to_home.py +++ b/src/modules/sbstudio/plugin/operators/return_to_home.py @@ -13,6 +13,7 @@ ) from sbstudio.plugin.model.formation import create_formation, get_markers_from_formation from sbstudio.plugin.model.safety_check import get_proximity_warning_threshold +from sbstudio.plugin.model.storyboard import get_storyboard from sbstudio.plugin.utils.evaluator import create_position_evaluator from .base import StoryboardOperator @@ -128,8 +129,9 @@ def draw(self, context): layout.prop(self, "use_smart_rth") def invoke(self, context, event): - storyboard = context.scene.skybrush.storyboard - self.start_frame = max(context.scene.frame_current, storyboard.frame_end) + self.start_frame = max( + context.scene.frame_current, get_storyboard(context).frame_end + ) return context.window_manager.invoke_props_dialog(self) def execute_on_storyboard(self, storyboard, entries, context): @@ -287,11 +289,8 @@ def _run(self, storyboard, *, context) -> bool: def _validate_start_frame(self, context: Context) -> bool: """Returns whether the return to home time chosen by the user is valid.""" - storyboard = context.scene.skybrush.storyboard - if storyboard.last_entry is not None: - last_frame = context.scene.skybrush.storyboard.frame_end - else: - last_frame = None + storyboard = get_storyboard(context) + last_frame = storyboard.frame_end if storyboard.last_entry is not None else None # TODO(ntamas): what if the last entry in the storyboard _is_ the # RTH, what shall we do then? Probably we should ignore it and look diff --git a/src/modules/sbstudio/plugin/operators/select_storyboard_entry.py b/src/modules/sbstudio/plugin/operators/select_storyboard_entry.py index ef63e7da..2dc4d1d3 100644 --- a/src/modules/sbstudio/plugin/operators/select_storyboard_entry.py +++ b/src/modules/sbstudio/plugin/operators/select_storyboard_entry.py @@ -1,3 +1,5 @@ +from sbstudio.plugin.model.storyboard import get_storyboard + from .base import StoryboardOperator __all__ = ("SelectStoryboardEntryForCurrentFrameOperator",) @@ -22,10 +24,7 @@ class SelectStoryboardEntryForCurrentFrameOperator(StoryboardOperator): @classmethod def poll(cls, context): - return ( - StoryboardOperator.poll(context) - and context.scene.skybrush.storyboard.entries - ) + return StoryboardOperator.poll(context) and get_storyboard(context).entries def execute_on_storyboard(self, storyboard, context): frame = context.scene.frame_current diff --git a/src/modules/sbstudio/plugin/operators/takeoff.py b/src/modules/sbstudio/plugin/operators/takeoff.py index bab5f2a0..afecbca6 100644 --- a/src/modules/sbstudio/plugin/operators/takeoff.py +++ b/src/modules/sbstudio/plugin/operators/takeoff.py @@ -9,7 +9,7 @@ from sbstudio.plugin.constants import Collections from sbstudio.plugin.model.formation import create_formation from sbstudio.plugin.model.safety_check import get_proximity_warning_threshold -from sbstudio.plugin.model.storyboard import Storyboard +from sbstudio.plugin.model.storyboard import get_storyboard, Storyboard from sbstudio.plugin.utils.evaluator import create_position_evaluator from .base import StoryboardOperator @@ -173,9 +173,10 @@ def _run(self, storyboard: Storyboard, *, context) -> bool: def _validate_start_frame(self, context: Context) -> bool: """Returns whether the takeoff time chosen by the user is valid.""" - storyboard = context.scene.skybrush.storyboard + storyboard = get_storyboard(context) # Note: we assume here that the first entry is the takeoff grid on ground if len(storyboard.entries) > 0: + assert storyboard.first_entry is not None frame = storyboard.first_entry.frame_end if self.start_frame < frame: self.report( @@ -187,6 +188,7 @@ def _validate_start_frame(self, context: Context) -> bool: ) return False if len(storyboard.entries) > 1: + assert storyboard.second_entry is not None frame = storyboard.second_entry.frame_start if frame is not None and self.start_frame >= frame: self.report( diff --git a/src/modules/sbstudio/plugin/panels/transition_editor.py b/src/modules/sbstudio/plugin/panels/transition_editor.py index b3db8543..a49993b7 100644 --- a/src/modules/sbstudio/plugin/panels/transition_editor.py +++ b/src/modules/sbstudio/plugin/panels/transition_editor.py @@ -2,7 +2,7 @@ from typing import List, Optional -from sbstudio.plugin.model.storyboard import Storyboard, StoryboardEntry +from sbstudio.plugin.model.storyboard import get_storyboard, Storyboard, StoryboardEntry from sbstudio.plugin.operators import ( CreateNewScheduleOverrideEntryOperator, RecalculateTransitionsOperator, @@ -36,8 +36,8 @@ def _get_entry(cls, storyboard: Storyboard) -> Optional[StoryboardEntry]: @classmethod def poll(cls, context) -> bool: - storyboard: Optional[Storyboard] = context.scene.skybrush.storyboard - if storyboard: + storyboard = get_storyboard(context) + if storyboard is not None: return cls._get_entry(storyboard) is not None else: return False @@ -53,10 +53,7 @@ def _get_recalculation_scope(cls) -> str: return "ALL" def draw(self, context): - storyboard: Optional[Storyboard] = context.scene.skybrush.storyboard - if not storyboard: - return - + storyboard = get_storyboard(context) entry = self._get_entry(storyboard) if entry is None: return diff --git a/src/modules/sbstudio/plugin/props/frame_range.py b/src/modules/sbstudio/plugin/props/frame_range.py index 9fb8dde5..38940915 100644 --- a/src/modules/sbstudio/plugin/props/frame_range.py +++ b/src/modules/sbstudio/plugin/props/frame_range.py @@ -3,6 +3,7 @@ from typing import Optional, Tuple +from sbstudio.plugin.model.storyboard import get_storyboard from sbstudio.plugin.utils import with_context __all__ = ("FrameRangeProperty",) @@ -48,14 +49,12 @@ def resolve_frame_range( return (context.scene.frame_preview_start, context.scene.frame_preview_end) elif range == "STORYBOARD": # Return the frame range covered by the storyboard - return ( - context.scene.skybrush.storyboard.frame_start, - context.scene.skybrush.storyboard.frame_end, - ) + storyboard = get_storyboard(context) + return (storyboard.frame_start, storyboard.frame_end) elif range == "AROUND_CURRENT_FRAME": # Return the frame range covered by the formation or transition around # the current frame - storyboard = context.scene.skybrush.storyboard + storyboard = get_storyboard(context) return storyboard.get_frame_range_of_formation_or_transition_at_frame( context.scene.frame_current )