From 9cebedc89fcabf509241d89b09bf39fb61cab94b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <45557362+qgallouedec@users.noreply.github.com>
Date: Fri, 5 May 2023 14:26:39 +0200
Subject: [PATCH] Fix Colab logger error (#1484)

* fix HumanOutputFormat

* update version

* update changelog

* TextIO annotation, TextIOBase isinstance

* update changelog

* test for HumanOutputFormat with custom TextIO

* rm extra test line

* Update tests/test_logger.py

Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>

---------

Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>
---
 docs/misc/changelog.rst            |  3 ++-
 stable_baselines3/common/logger.py |  4 +--
 stable_baselines3/version.txt      |  2 +-
 tests/test_logger.py               | 39 ++++++++++++++++++++++++++++++
 4 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
index 60750c2..096648a 100644
--- a/docs/misc/changelog.rst
+++ b/docs/misc/changelog.rst
@@ -3,7 +3,7 @@
 Changelog
 ==========
 
-Release 2.0.0a7 (WIP)
+Release 2.0.0a8 (WIP)
 --------------------------
 
 **Gymnasium support**
@@ -21,6 +21,7 @@ Breaking Changes:
 - Removed deprecated ``stack_observation_space`` method of ``StackedObservations``
 - Renamed environment output observations in ``evaluate_policy`` to prevent shadowing the input observations during callbacks (@npit)
 - Upgraded wrappers and custom environment to Gymnasium
+- Refined the ``HumanOutputFormat`` file check: now it verifies if the object is an instance of ``io.TextIOBase`` instead of only checking for the presence of a ``write`` method.
 
 New Features:
 ^^^^^^^^^^^^^
diff --git a/stable_baselines3/common/logger.py b/stable_baselines3/common/logger.py
index f1aadab..8a6facf 100644
--- a/stable_baselines3/common/logger.py
+++ b/stable_baselines3/common/logger.py
@@ -5,7 +5,7 @@ import sys
 import tempfile
 import warnings
 from collections import defaultdict
-from io import TextIOWrapper
+from io import TextIOBase
 from typing import Any, Dict, List, Mapping, Optional, Sequence, TextIO, Tuple, Union
 
 import numpy as np
@@ -164,7 +164,7 @@ class HumanOutputFormat(KVWriter, SeqWriter):
         if isinstance(filename_or_file, str):
             self.file = open(filename_or_file, "w")
             self.own_file = True
-        elif isinstance(filename_or_file, TextIOWrapper):  # equivalent to `isinstance(..., TextIO)` (not supported)
+        elif isinstance(filename_or_file, TextIOBase):
             self.file = filename_or_file
             self.own_file = False
         else:
diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt
index 4b23e04..55a349a 100644
--- a/stable_baselines3/version.txt
+++ b/stable_baselines3/version.txt
@@ -1 +1 @@
-2.0.0a7
+2.0.0a8
diff --git a/tests/test_logger.py b/tests/test_logger.py
index c33bd4c..7a1c389 100644
--- a/tests/test_logger.py
+++ b/tests/test_logger.py
@@ -2,6 +2,7 @@ import importlib.util
 import os
 import sys
 import time
+from io import TextIOBase
 from typing import Sequence
 from unittest import mock
 
@@ -434,3 +435,41 @@ def test_ep_buffers_stats_window_size(algo, stats_window_size):
     model.learn(total_timesteps=10)
     assert model.ep_info_buffer.maxlen == stats_window_size
     assert model.ep_success_buffer.maxlen == stats_window_size
+
+
+def test_human_output_format_custom_test_io():
+    class DummyTextIO(TextIOBase):
+        def __init__(self) -> None:
+            super().__init__()
+            self.lines = [[]]
+
+        def write(self, t: str) -> int:
+            self.lines[-1].append(t)
+
+        def flush(self) -> None:
+            self.lines.append([])
+
+        def close(self) -> None:
+            pass
+
+        def get_printed(self) -> str:
+            return "\n".join(["".join(line) for line in self.lines])
+
+    dummy_text_io = DummyTextIO()
+    output = HumanOutputFormat(dummy_text_io)
+    output.write({"key1": "value1", "key2": 42}, {"key1": None, "key2": None})
+    output.write({"key1": "value2", "key2": 43}, {"key1": None, "key2": None})
+    printed = dummy_text_io.get_printed()
+    desired_printed = """-----------------
+| key1 | value1 |
+| key2 | 42     |
+-----------------
+
+-----------------
+| key1 | value2 |
+| key2 | 43     |
+-----------------
+
+"""
+
+    assert printed == desired_printed