diff --git a/.gitignore b/.gitignore
index 56e4b6d4b43e9b0c1becd55c01a26e1855d58cda..da2082c088ce09c27045e816bc0e4d9b4d00d910 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,7 +8,7 @@ objs
 # Python items
 cython_debug/
 python_build/
-python_format_venv/
+yapf_virtual_environment/
 python_pylint_venv/
 .coverage*
 .eggs
diff --git a/tools/distrib/yapf_code.sh b/tools/distrib/yapf_code.sh
index 007b14810eb9e01cdac3d9c9ca01e43042dafc2d..f28a1ce8baeb756444487adfa66e4d8b40794210 100755
--- a/tools/distrib/yapf_code.sh
+++ b/tools/distrib/yapf_code.sh
@@ -31,31 +31,48 @@
 set -ex
 
 # change to root directory
-cd $(dirname $0)/../..
+cd "$(dirname "${0}")/../.."
 
-DIRS=src/python
-EXCLUSIONS='src/python/grpcio/grpc_*.py src/python/grpcio_health_checking/grpc_*.py src/python/grpcio_reflection/grpc_*.py src/python/grpcio_tests/grpc_*.py'
+DIRS=(
+    'src/python'
+)
+EXCLUSIONS=(
+    'grpcio/grpc_*.py'
+    'grpcio_health_checking/grpc_*.py'
+    'grpcio_reflection/grpc_*.py'
+    'grpcio_tests/grpc_*.py'
+)
 
-VIRTUALENV=python_format_venv
+VIRTUALENV=yapf_virtual_environment
 
 virtualenv $VIRTUALENV
-PYTHON=`realpath $VIRTUALENV/bin/python`
-$PYTHON -m pip install futures
+PYTHON=$(realpath "${VIRTUALENV}/bin/python")
+$PYTHON -m pip install --upgrade pip
+$PYTHON -m pip install --upgrade futures
 $PYTHON -m pip install yapf==0.16.0
 
-exclusion_args=""
-for exclusion in $EXCLUSIONS; do
-  exclusion_args="$exclusion_args --exclude $exclusion"
-done
+yapf() {
+    local exclusion exclusion_args=()
+    for exclusion in "${EXCLUSIONS[@]}"; do
+        exclusion_args+=( "--exclude" "$1/${exclusion}" )
+    done
+    $PYTHON -m yapf -i -r --style=setup.cfg -p "${exclusion_args[@]}" "${1}"
+}
 
-script_result=0
-for dir in $DIRS; do
-  tempdir=`mktemp -d`
-  cp -RT $dir $tempdir
-  $PYTHON -m yapf -i -r -p $exclusion_args $dir
-  if ! diff -r $dir $tempdir; then
-    script_result=1
-  fi
-  rm -rf $tempdir
-done
-exit $script_result
+if [[ -z "${TEST}" ]]; then
+    for dir in "${DIRS[@]}"; do
+	yapf "${dir}"
+    done
+else
+    ok=yes
+    for dir in "${DIRS[@]}"; do
+	tempdir=$(mktemp -d)
+	cp -RT "${dir}" "${tempdir}"
+	yapf "${tempdir}"
+	diff -ru "${dir}" "${tempdir}" || ok=no
+	rm -rf "${tempdir}"
+    done
+    if [[ ${ok} == no ]]; then
+	false
+    fi
+fi