Clean up code

server/books/models/snapshot.py

@@ -86,7 +86,6 @@ class SnapshotManager(models.Manager):
                     order=objective.order
                 )
 
-
         return snapshot
 
 
@@ -121,3 +120,33 @@ class Snapshot(models.Model):
 
     def __str__(self):
         return f'Snapshot {self.id}'
+
+    def reset(self, user, selected_class):
+        for chapter in Chapter.get_by_parent(self.module):
+            qs = ContentBlock.get_by_parent(chapter)
+            without_owner = Q(owner__isnull=True)
+            no_snapshot = Q(contentblocksnapshot__isnull=True)
+            owner_user = Q(owner=user)
+            for content_block in qs.filter(without_owner & no_snapshot):
+                content_block.hidden_for.remove(selected_class)
+            for content_block in qs.filter(owner_user):
+                content_block.visible_for.remove(selected_class)
+
+    def apply(self, user, selected_class):
+        for content_block in self.hidden_content_blocks.all():
+            content_block.hidden_for.add(selected_class)
+        for custom_content_block in self.custom_content_blocks.all():
+            custom_content_block.to_regular_content_block(owner=user, school_class=selected_class)
+        for chapter_snapshot in self.chapters.through.objects.all():
+            chapter = chapter_snapshot.chapter
+            if chapter_snapshot.title_hidden:
+                chapter.title_hidden_for.add(selected_class)
+            if chapter_snapshot.description_hidden:
+                chapter.description_hidden_for.add(selected_class)
+        for objective_group_snapshot in self.objective_groups.through.objects.all():
+            if objective_group_snapshot.hidden:
+                objective_group_snapshot.objective_group.hidden_for.add(selected_class)
+        for objective in self.hidden_objectives.all():
+            objective.hidden_for.add(selected_class)
+        for custom_objective in self.custom_objectives.all():
+            custom_objective.to_regular_objective(owner=user, school_class=selected_class)

server/books/schema/mutations/snapshot.py

@@ -45,36 +45,14 @@ class ApplySnapshot(relay.ClientIDMutation):
         user = info.context.user
         selected_class_id = args.get('selected_class')
         selected_class = get_object(SchoolClass, selected_class_id)
-        # reset everything
+        # permission check
-        for chapter in Chapter.get_by_parent(snapshot.module):
-            cb_qs = ContentBlock.get_by_parent(chapter)
-            without_owner = Q(owner__isnull=True)
-            no_snapshot = Q(contentblocksnapshot__isnull=True)
-            owner_user = Q(owner=user)
-            for cb in cb_qs.filter(without_owner & no_snapshot):
-                cb.hidden_for.remove(selected_class)
-            for cb in cb_qs.filter(owner_user):
-                cb.visible_for.remove(selected_class)
-        # apply snapshot
         if not selected_class.users.filter(username=user.username).exists() or not user.is_teacher():
             raise PermissionError('Not allowed')
-        for content_block in snapshot.hidden_content_blocks.all():
+
-            content_block.hidden_for.add(selected_class)
+        # reset everything
-        for custom_content_block in snapshot.custom_content_blocks.all():
+        snapshot.reset(user, selected_class)
-            custom_content_block.to_regular_content_block(owner=user, school_class=selected_class)
+        # apply snapshot
-        for chapter_snapshot in snapshot.chapters.through.objects.all():
+        snapshot.apply(user, selected_class)
-            chapter = chapter_snapshot.chapter
-            if chapter_snapshot.title_hidden:
-                chapter.title_hidden_for.add(selected_class)
-            if chapter_snapshot.description_hidden:
-                chapter.description_hidden_for.add(selected_class)
-        for objective_group_snapshot in snapshot.objective_groups.through.objects.all():
-            if objective_group_snapshot.hidden:
-                objective_group_snapshot.objective_group.hidden_for.add(selected_class)
-        for objective in snapshot.hidden_objectives.all():
-            objective.hidden_for.add(selected_class)
-        for custom_objective in snapshot.custom_objectives.all():
-            custom_objective.to_regular_objective(owner=user, school_class=selected_class)
         return cls(success=True, module=snapshot.module)