From 1806b2890b7546fd3d1955f8b4cb07279f22ff54 Mon Sep 17 00:00:00 2001
From: Fedor Baart <fedor.baart@deltares.nl>
Date: Sat, 23 Apr 2022 07:19:09 +0200
Subject: [PATCH] Feature/more logging (#191)

* document new feature

* test for maxLogLength

* add maximum log length parameter

* remove unused variable

* consistent implementation
---
 README.md                                      |  1 +
 kaggle_environments/agent.py                   | 17 ++++++++++++-----
 kaggle_environments/core.py                    | 14 ++++++++++----
 .../envs/connectx/test_connectx.py             | 18 ++++++++++++++++++
 4 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 6d905b69..624da332 100644
--- a/README.md
+++ b/README.md
@@ -200,6 +200,7 @@ There are two types of configuration: Defaults applying to every environment and
 | agentTimeout | Maximum runtime (seconds) to initialize an agent.               |
 | actTimeout   | Maximum runtime (seconds) to obtain an action from an agent.    |
 | runTimeout   | Maximum runtime (seconds) of an episode (not necessarily DONE). |
+| maxLogLength | Maximum log length (number of characters, `None` -> no limit)   |
 
 ```python
 env = make("connectx", configuration={
diff --git a/kaggle_environments/agent.py b/kaggle_environments/agent.py
index 29e0ded2..b40f351b 100644
--- a/kaggle_environments/agent.py
+++ b/kaggle_environments/agent.py
@@ -46,7 +46,7 @@ def get_last_callable(raw, fallback=None, path=None):
         # append exec_dir so that way python agents can import other files
         exec_dir = os.path.dirname(path)
         sys.path.append(exec_dir)
-        
+
         exec(code_object, env)
         sys.path.pop()
         sys.stdout = orig_out
@@ -160,10 +160,17 @@ def act(self, observation):
             except Exception as e:
                 traceback.print_exc(file=err_buffer)
                 action = e
-            # Allow up to 1k log characters per step which is ~1MB per 600 step episode
-            max_log_length = 1024
-            out = out_buffer.getvalue()[0:max_log_length]
-            err = err_buffer.getvalue()[0:max_log_length]
+
+            out = out_buffer.getvalue()
+            err = err_buffer.getvalue()
+            # Get the maximum log length
+            # Allow up to 1k (default) log characters per step which is ~1MB per 600 step episode
+            max_log_length = self.configuration.get('maxLogLength', 1024)
+
+            # truncate if max_log_length is set to None, do not truncate
+            if max_log_length is not None:
+                out = out[0:max_log_length]
+                err = err[0:max_log_length]
 
         duration = perf_counter() - start
         log = {
diff --git a/kaggle_environments/core.py b/kaggle_environments/core.py
index c6dcc3b6..2ee46577 100644
--- a/kaggle_environments/core.py
+++ b/kaggle_environments/core.py
@@ -600,14 +600,20 @@ def __run_interpreter(self, state, logs):
                     # Reraise e to ensure that the program exits
                     raise e
                 finally:
-                    # Allow up to 1k log characters per step which is ~1MB per 600 step episode
-                    max_log_length = 1024
                     out = out_buffer.getvalue()
                     err = err_buffer.getvalue()
+
+                    # strip if needed
+                    # Allow up to 1k (default) log characters per step which is ~1MB per 600 step episode
+                    max_log_length = self.configuration.get("maxLogLength", 1024)
+                    if max_log_length is not None:
+                        out = out[0:max_log_length]
+                        err = err[0:max_log_length]
+
                     if out or err:
                         logs.append({
-                            "stdout": out[0:max_log_length],
-                            "stderr": err[0:max_log_length]
+                            "stdout": out,
+                            "stderr": err
                         })
         finally:
             if out:
diff --git a/kaggle_environments/envs/connectx/test_connectx.py b/kaggle_environments/envs/connectx/test_connectx.py
index e59a2ce4..81045ba0 100644
--- a/kaggle_environments/envs/connectx/test_connectx.py
+++ b/kaggle_environments/envs/connectx/test_connectx.py
@@ -257,3 +257,21 @@ def test_can_evaluate():
     rewards = evaluate("connectx", ["random", "random"], num_episodes=2)
     assert (rewards[0][0] + rewards[0][1] ==
             0) and rewards[1][0] + rewards[1][1] == 0
+
+
+def test_max_log_length():
+    def custom1():
+        # Write 20X to stdtout, we should strip to 10
+        print('X' * 20)
+        return 1
+
+    def custom2():
+        return 2
+    before_each(
+        # here we strip log to length 10
+        configuration={"rows": 4, "columns": 5, "inarow": 3, "maxLogLength": 10},
+    )
+    env.run([custom1, custom2])
+    last_log = env.logs[-1][0]['stdout']
+    assert env.configuration.maxLogLength == 10, "max log length should be set to 10"
+    assert len(last_log.strip()) == 10, "max log length should be 10 (+ newline, which we stripped)"