From d53d0c21c2a193a03c2cb3d12b354bb191292df8 Mon Sep 17 00:00:00 2001
From: Jon Herron <jon.herron@yahoo.com>
Date: Wed, 22 Feb 2023 13:13:28 -0500
Subject: [PATCH] Provide info about keys in task data and python env

---
 .../scripts/get_data_sizes.py                 | 21 ++++++++++++-------
 .../services/process_instance_processor.py    | 15 +++++++++----
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/src/spiffworkflow_backend/scripts/get_data_sizes.py b/src/spiffworkflow_backend/scripts/get_data_sizes.py
index 94dd2fd6..553d85c5 100644
--- a/src/spiffworkflow_backend/scripts/get_data_sizes.py
+++ b/src/spiffworkflow_backend/scripts/get_data_sizes.py
@@ -32,12 +32,17 @@ class GetDataSizes(Script):
         **kwargs: Any
     ) -> Any:
         """Run."""
-        task = script_attributes_context.task
-        cumulative_task_data_size = ProcessInstanceProcessor.get_task_data_size(
-            task.workflow
-        )
-        python_env_size = ProcessInstanceProcessor.get_python_env_size(task.workflow)
-        return {
-            "cumulative_task_data_size": cumulative_task_data_size,
-            "python_env_size": python_env_size,
+        workflow = script_attributes_context.task.workflow
+        task_data_size = ProcessInstanceProcessor.get_task_data_size(workflow)
+        task_data_keys_by_task = {
+            t.task_spec.name: sorted(t.data.keys())
+            for t in ProcessInstanceProcessor.get_tasks_with_data(workflow)
+        }
+        python_env_size = ProcessInstanceProcessor.get_python_env_size(workflow)
+        python_env_keys = workflow.script_engine.environment.user_defined_state().keys()
+        return {
+            "python_env_size": python_env_size,
+            "python_env_keys": sorted(python_env_keys),
+            "task_data_size": task_data_size,
+            "task_data_keys_by_task": task_data_keys_by_task,
         }
diff --git a/src/spiffworkflow_backend/services/process_instance_processor.py b/src/spiffworkflow_backend/services/process_instance_processor.py
index 5b728cd6..498ebadf 100644
--- a/src/spiffworkflow_backend/services/process_instance_processor.py
+++ b/src/spiffworkflow_backend/services/process_instance_processor.py
@@ -1607,14 +1607,21 @@ class ProcessInstanceProcessor:
         except WorkflowTaskException as we:
             raise ApiError.from_workflow_exception("task_error", str(we), we) from we
 
+    @classmethod
+    def get_tasks_with_data(cls, bpmn_process_instance: BpmnWorkflow) -> List[SpiffTask]:
+        return [
+            task
+            for task in bpmn_process_instance.get_tasks(TaskState.FINISHED_MASK)
+            if len(task.data) > 0
+        ]
+
     @classmethod
     def get_task_data_size(cls, bpmn_process_instance: BpmnWorkflow) -> int:
-        tasks_to_check = bpmn_process_instance.get_tasks(TaskState.FINISHED_MASK)
-        task_data = [task.data for task in tasks_to_check]
-        task_data_to_check = list(filter(len, task_data))
+        tasks_with_data = cls.get_tasks_with_data(bpmn_process_instance)
+        all_task_data = [task.data for task in tasks_with_data]
 
         try:
-            return len(json.dumps(task_data_to_check))
+            return len(json.dumps(all_task_data))
         except Exception:
             return 0