Fix for python 3

.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

python-package/xgboost/plotting.py

@@ -8,11 +8,7 @@ import re
 import numpy as np
 from .core import Booster
 
-try:
+from io import BytesIO
-    from StringIO import StringIO
-except ImportError:
-    from io import StringIO
-
 
 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)

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

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

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