From 75200671c4d52de7d2f7ae33e19ed3c051bdc592 Mon Sep 17 00:00:00 2001
From: Peter Zhokhov <peterz@openai.com>
Date: Fri, 3 May 2019 16:29:10 -0700
Subject: [PATCH] fix tests - add matplotlib to setup_requires, put mpi4py
 import in try-except

---
 baselines/common/test_mpi_util.py | 8 ++++++--
 setup.py                          | 3 ++-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/baselines/common/test_mpi_util.py b/baselines/common/test_mpi_util.py
index 5d1e7a5..ce70fde 100644
--- a/baselines/common/test_mpi_util.py
+++ b/baselines/common/test_mpi_util.py
@@ -1,7 +1,11 @@
 from baselines.common import mpi_util
-from mpi4py import MPI
 from baselines import logger
 from baselines.common.tests.test_with_mpi import with_mpi
+try:
+    from mpi4py import MPI
+except ImportError:
+    MPI = None
+    
 
 @with_mpi()
 def test_mpi_weighted_mean():
@@ -23,4 +27,4 @@ def test_mpi_weighted_mean():
                 logger.logkv_mean(name, val)
         d2 = logger.dumpkvs()
         if comm.rank == 0:
-            assert d2 == correctval
\ No newline at end of file
+            assert d2 == correctval
diff --git a/setup.py b/setup.py
index 97c56a6..46a1819 100644
--- a/setup.py
+++ b/setup.py
@@ -12,7 +12,8 @@ extras = {
         'filelock',
         'pytest',
         'pytest-forked',
-        'atari-py'
+        'atari-py',
+        'matplotlib',
     ],
     'mpi': [
         'mpi4py'