From 47f0ee1c01e13924c4140cf955bb7af09f5a154d Mon Sep 17 00:00:00 2001
From: Jakob Langdal <jakob.langdal@alexandra.dk>
Date: Mon, 17 Apr 2023 21:34:49 +0000
Subject: [PATCH] Add option to exclude pickled result

---
 optimizerapi/optimizer.py |  5 ++++-
 tests/test_optimizer.py   | 26 ++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/optimizerapi/optimizer.py b/optimizerapi/optimizer.py
index e5700e7..c0a81db 100644
--- a/optimizerapi/optimizer.py
+++ b/optimizerapi/optimizer.py
@@ -179,6 +179,8 @@ def process_result(result, optimizer, dimensions, cfg, extras, data, space):
 
     objective_pars = extras.get("objectivePars", "result")
 
+    pickle_model = json.loads(extras.get("includeModel", "true").lower())
+
     # In the following section details that should be reported to
     # clients should go into the "resultDetails" dictionary and plots
     # go into the "plots" list (this is handled by calling the "addPlot" function)
@@ -220,7 +222,8 @@ def process_result(result, optimizer, dimensions, cfg, extras, data, space):
                 plot_Pareto(optimizer)
                 add_plot(plots, "pareto")
 
-    result_details["pickled"] = pickleToString(result, get_crypto())
+    if pickle_model:
+        result_details["pickled"] = pickleToString(result, get_crypto())
 
     add_version_info(result_details["extras"])
 
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 32b487d..628fee4 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -105,3 +105,29 @@ def test_can_accept_multi_objective_data():
     validateResult(result)
     assert len(result["result"]["models"]) > 1
     assert len(result["plots"]) == 5
+
+
+def test_deselecting_pickled_model():
+    # If includeModel is false, pickled data should not be included in result
+    result = optimizer.run(
+        body={
+            "data": sampleData,
+            "optimizerConfig": sampleConfig,
+            "extras": {"includeModel": "false"},
+        }
+    )
+    assert "pickled" in result["result"]
+    assert len(result["result"]["pickled"]) == 0
+
+
+def test_selecting_pickled_model():
+    # If includeModel is true, pickled data should be included in result
+    result = optimizer.run(
+        body={
+            "data": sampleData,
+            "optimizerConfig": sampleConfig,
+            "extras": {"includeModel": "true"},
+        }
+    )
+    assert "pickled" in result["result"]
+    assert len(result["result"]["pickled"]) > 0