Deal with yaml import if need be

1 files changed, 9 insertions, 1 deletions

diff --git a/hooks/hooks.py b/hooks/hooks.py
index 95a67d6..941caed 100755
--- a/hooks/hooks.py
+++ b/hooks/hooks.py
@@ -13,7 +13,15 @@ import socket
 import subprocess
 import sys
 import time
-import yaml
+
+try:
+ import yaml # flake8: noqa
+except ImportError:
+ if sys.version_info.major == 2:
+ subprocess.check_call(['apt-get', 'install', '-y', 'python-yaml'])
+ else:
+ subprocess.check_call(['apt-get', 'install', '-y', 'python3-yaml'])
+ import yaml # flake8: noqa
 
 from os import chmod
 from os import remove