From a898a8f1f090d84efc627c19f43fdfc34f797d25 Mon Sep 17 00:00:00 2001
From: Jesse Hellemn <hellemn@fb.com>
Date: Wed, 27 Jun 2018 13:01:51 -0700
Subject: [PATCH] Adding pyyaml to mac and windows builds

Summary: Closes https://github.com/pytorch/pytorch/pull/8851

Reviewed By: mingzhe09088

Differential Revision: D8666075

Pulled By: pjh5

fbshipit-source-id: a3fdc9f9801f814b1e4010bd20ba51afbb048a1d
---
 scripts/build_local.sh    | 6 ++++++
 scripts/build_windows.bat | 3 +++
 2 files changed, 9 insertions(+)

diff --git a/scripts/build_local.sh b/scripts/build_local.sh
index 23e1394dfff..ea684de04d6 100755
--- a/scripts/build_local.sh
+++ b/scripts/build_local.sh
@@ -31,6 +31,12 @@ if [ -n "${USE_ANACONDA}" ]; then
   export CONDA_INSTALL_LOCALLY=1
   "${ROOT_DIR}/scripts/build_anaconda.sh" "$@"
 else
+  # Make sure that pyyaml is installed for the codegen of building Aten to work
+  if [[ -n "$(python -c 'import yaml' 2>&1)" ]]; then
+    echo "Installing pyyaml with pip at $(which pip)"
+    pip install --user pyyaml
+  fi
+
   # Build protobuf compiler from third_party if configured to do so
   if [ -n "${USE_HOST_PROTOC:-}" ]; then
     echo "USE_HOST_PROTOC is set; building protoc before building Caffe2..."
diff --git a/scripts/build_windows.bat b/scripts/build_windows.bat
index 185d43ac0a8..c4b8a38fafa 100644
--- a/scripts/build_windows.bat
+++ b/scripts/build_windows.bat
@@ -37,6 +37,9 @@ if NOT DEFINED CMAKE_GENERATOR (
   )
 )
 
+:: Install pyyaml for Aten codegen
+pip install pyyaml
+
 echo CAFFE2_ROOT=%CAFFE2_ROOT%
 echo CMAKE_GENERATOR=%CMAKE_GENERATOR%
 echo CMAKE_BUILD_TYPE=%CMAKE_BUILD_TYPE%