From 6217934c99ce7defc5ebe8c4070ae3596c5f1f18 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@stuvel.eu>
Date: Wed, 9 Nov 2016 17:12:18 +0100
Subject: [PATCH] Task types: also take parent node type into account

---
 attract/tasks/routes.py | 46 ++++++++++++++++++++++++++++++++++-------
 1 file changed, 39 insertions(+), 7 deletions(-)

diff --git a/attract/tasks/routes.py b/attract/tasks/routes.py
index 1c142f4..6a45bc8 100644
--- a/attract/tasks/routes.py
+++ b/attract/tasks/routes.py
@@ -13,6 +13,7 @@ import pillar.web.subquery
 
 from attract.routes import attract_project_view
 from attract.node_types.task import node_type_task
+from attract.node_types.shot import node_type_shot
 from attract import current_attract, ROLES_REQUIRED_TO_VIEW_ITEMS, EXTENSION_NAME
 
 blueprint = Blueprint('attract.tasks', __name__, url_prefix='/tasks')
@@ -69,13 +70,7 @@ def view_task(project, attract_props, task_id):
 
     # Figure out which task types are available, defaulting to the shot task types.
     context = request.args.get('context', None) or 'shot'
-    ctx_node_type_name = '%s_%s' % (EXTENSION_NAME, context)
-    try:
-        task_types = attract_props['task_types'][ctx_node_type_name]
-    except KeyError:
-        log.warning('Project %s does not have an Attract task type definition for %s',
-                    project['_id'], ctx_node_type_name)
-        task_types = []
+    task_types = task_types_given_context(project, attract_props, context, task)
 
     if task.properties.due_date:
         task.properties.due_date = parser.parse('%s' % task.properties.due_date)
@@ -102,6 +97,43 @@ def view_task(project, attract_props, task_id):
                            attract_context=request.args.get('context'))
 
 
+def task_types_given_context(project, attract_props, page_context, task):
+    """Returns a list of task types, given the page context and/or task parent type."""
+
+    # If we're in an explicit shot/asset context, just use that.
+    if page_context in {'shot', 'asset'}:
+        ctx_node_type_name = '%s_%s' % (EXTENSION_NAME, page_context)
+        try:
+            return attract_props['task_types'][ctx_node_type_name]
+        except KeyError:
+            log.warning('Project %s does not have an Attract task type definition for %s',
+                        project['_id'], ctx_node_type_name)
+            # Fall through to the case below.
+
+    # If we're not in such a context, we need to inspect the parent node type (if any).
+    if task.parent:
+        api = pillar_api()
+        parent = pillarsdk.Node.find(task.parent, {'projection': {'node_type': 1}}, api=api)
+        if parent:
+            try:
+                return attract_props['task_types'][parent['node_type']]
+            except KeyError:
+                log.warning('Project %s does not have an Attract task type definition for %s',
+                            project['_id'], parent['node_type'])
+                # Fall through to the fallback case below.
+
+    # Just fall back to shot task types
+    try:
+        return attract_props['task_types'][node_type_shot['name']]
+    except KeyError:
+        log.warning('Project %s does not have an Attract task type definition for %s',
+                    project['_id'], parent['node_type'])
+        # Fall through to the fallback case below.
+
+    # Fallback in case of total failure.
+    return []
+
+
 @perproject_blueprint.route('/<task_id>', methods=['POST'])
 @attract_project_view()
 def save(project, task_id):