From 00702dc39b18ec7b74d3b7e8bf1a55b2dee75126 Mon Sep 17 00:00:00 2001
From: sinhrks <sinhrks@gmail.com>
Date: Sun, 23 Aug 2015 13:37:29 +0900
Subject: [PATCH] Fix for python 3

---
 .travis.yml                        |  4 +---
 python-package/xgboost/plotting.py |  8 ++------
 scripts/travis_osx_install.sh      |  8 +++++++-
 scripts/travis_script.sh           | 21 ++++++++++++++++++++-
 tests/python/test_basic.py         |  3 +++
 5 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 10f217038..5e10c3360 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -12,6 +12,7 @@ env:
     - TASK=lint LINT_LANG=python
     - TASK=R-package CXX=g++
     - TASK=python-package CXX=g++
+    - TASK=python-package3 CXX=g++
     - TASK=java-package CXX=g++
     - TASK=build CXX=g++
     - TASK=build-with-dmlc CXX=g++
@@ -29,9 +30,6 @@ addons:
       - wget
       - libcurl4-openssl-dev
       - unzip
-      - python-numpy
-      - python-scipy
-      - python-nose
 
 before_install:
   - scripts/travis_osx_install.sh
diff --git a/python-package/xgboost/plotting.py b/python-package/xgboost/plotting.py
index 7c34a11f1..9c9b2a97d 100644
--- a/python-package/xgboost/plotting.py
+++ b/python-package/xgboost/plotting.py
@@ -8,11 +8,7 @@ import re
 import numpy as np
 from .core import Booster
 
-try:
-    from StringIO import StringIO
-except ImportError:
-    from io import StringIO
-
+from io import BytesIO
 
 def plot_importance(booster, ax=None, height=0.2,
                     xlim=None, title='Feature importance',
@@ -217,7 +213,7 @@ def plot_tree(booster, num_trees=0, rankdir='UT', ax=None, **kwargs):
 
     g = to_graphviz(booster, num_trees=num_trees, rankdir=rankdir, **kwargs)
 
-    s = StringIO()
+    s = BytesIO()
     s.write(g.pipe(format='png'))
     s.seek(0)
     img = image.imread(s)
diff --git a/scripts/travis_osx_install.sh b/scripts/travis_osx_install.sh
index 9f3f6e831..adc620a52 100755
--- a/scripts/travis_osx_install.sh
+++ b/scripts/travis_osx_install.sh
@@ -9,5 +9,11 @@ brew update
 if [ ${TASK} == "python-package" ]; then
     brew install python git graphviz
     easy_install pip
-    pip install numpy scipy nose
+    pip install numpy scipy matplotlib nose
+fi
+
+if [ ${TASK} == "python-package3" ]; then
+    brew install python3 git graphviz
+    sudo pip3 install --upgrade setuptools
+    pip3 install numpy scipy matplotlib nose graphviz
 fi
diff --git a/scripts/travis_script.sh b/scripts/travis_script.sh
index d58eaa5de..c5708b0c8 100755
--- a/scripts/travis_script.sh
+++ b/scripts/travis_script.sh
@@ -35,11 +35,30 @@ fi
 
 if [ ${TASK} == "python-package" ]; then
     sudo apt-get install graphviz
-    sudo pip install matplotlib graphviz
+    sudo apt-get install python-numpy python-scipy python-matplotlib python-nose
+    sudo python -m pip install graphviz
     make all CXX=${CXX} || exit -1
     nosetests tests/python || exit -1
 fi
 
+if [ ${TASK} == "python-package3" ]; then
+    sudo apt-get install graphviz
+    # python3-matplotlib is unavailale on Ubuntu 12.04
+    sudo apt-get install python3-dev
+    sudo apt-get install python3-numpy python3-scipy python3-nose python3-setuptools
+
+    make all CXX=${CXX} || exit -1
+
+    if [ ${TRAVIS_OS_NAME} != "osx" ]; then
+        sudo easy_install3 pip
+        sudo easy_install3 -U distribute
+        sudo pip install graphviz matplotlib
+        nosetests3 tests/python || exit -1
+    else
+        nosetests tests/python || exit -1
+    fi
+fi
+
 # only test java under linux for now
 if [ ${TASK} == "java-package" ]; then
     if [ ${TRAVIS_OS_NAME} != "osx" ]; then
diff --git a/tests/python/test_basic.py b/tests/python/test_basic.py
index 2ed1cc462..93ebaa7fd 100644
--- a/tests/python/test_basic.py
+++ b/tests/python/test_basic.py
@@ -33,6 +33,9 @@ def test_plotting():
     bst2 = xgb.Booster(model_file='xgb.model')
     # plotting
 
+    import matplotlib
+    matplotlib.use('Agg')
+
     from matplotlib.axes import Axes
     from graphviz import Digraph