From c92fb73e91aeb32a5507c9f00f19969c1fa775d5 Mon Sep 17 00:00:00 2001
From: "Maxence G. de Montauzan" <maxence@gdemontauzan.fr>
Date: Sun, 20 Dec 2020 18:39:04 +0100
Subject: [PATCH] (send_data) Improve script output

---
 send_data.py | 31 +++++++++++++++++++++----------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/send_data.py b/send_data.py
index b40ed6f..35c573d 100644
--- a/send_data.py
+++ b/send_data.py
@@ -10,6 +10,17 @@ import json
 import time
 import requests
 
+class bcolors:
+    HEADER = '\033[95m'
+    OKBLUE = '\033[94m'
+    OKCYAN = '\033[96m'
+    OKGREEN = '\033[92m'
+    WARNING = '\033[93m'
+    FAIL = '\033[91m'
+    ENDC = '\033[0m'
+    BOLD = '\033[1m'
+    UNDERLINE = '\033[4m'
+
 # Default file names
 DEFAULT_SONG_FILE = 'es-songs.json'
 DEFAULT_ALBUM_FILE = 'es-albums.json'
@@ -231,7 +242,7 @@ def send_data(file, quiet=False):
         print(res.text)
     else:
         if not quiet:
-            print("File '{}' sent to Elasticsearch!".format(file.name))
+            print(bcolors.OKGREEN + "File '{}' sent to Elasticsearch!".format(file.name) + bcolors.ENDC)
 
 def delete_index(index_name, quiet=False):
     """
@@ -242,9 +253,9 @@ def delete_index(index_name, quiet=False):
     res = requests.delete(url=ELASTICSEARCH_URL + index_name)
     if res.status_code == 200:
         if not quiet:
-            print("Deleted!")
+            print(bcolors.OKGREEN + "Index deleted!" + bcolors.ENDC)
     else:
-        print("An error occured")
+        print(bcolors.FAIL + "An error occured" + bcolors.ENDC)
         if res.json()['error']['type'] == 'index_not_found_exception':
             print("Index '{}' doesn't exist and can't be deleted".format(index_name))
         else:
@@ -261,11 +272,11 @@ def put_mapping(index_name, mapping_file, quiet=False):
                        data=mapping_file,
                        headers={'Content-Type': 'application/json'})
     if res.status_code != 200:
-        print("An error occured")
-        print(res.text)
+        print(bcolors.FAIL + "An error occured")
+        print(res.text + bcolors.ENDC)
     else:
         if not quiet:
-            print("File '{}' sent to Elasticsearch!".format(mapping_file.name))
+            print(bcolors.OKGREEN + "Mapping sent" + bcolors.ENDC)
 
     put_setting(index_name, 0, quiet)
 
@@ -292,8 +303,8 @@ def check_all_data_is_saved(data_file, index_name, quiet=False):
                        data=json.dumps(payload),
                        headers={'Content-Type': 'application/x-ndjson'})
     if res.status_code != 200:
-        print("An error occured")
-        print(res.text)
+        print(bcolors.FAIL + "An error occured")
+        print(res.text + bcolors.ENDC)
 
     els_nb_doc = res.json()['hits']['total']['value']
 
@@ -301,9 +312,9 @@ def check_all_data_is_saved(data_file, index_name, quiet=False):
         print("\tFound: {} documents in index '{}' in ELS".format(els_nb_doc, index_name))
 
     if file_nb_line != els_nb_doc:
-        print('Look out! Not all the data has been found in ELS')
+        print(bcolors.WARNING + 'Look out! Not all the data has been found in ELS' + bcolors.ENDC)
     elif not quiet:
-        print('All data is in ELS, it\'s ok')
+        print(bcolors.OKGREEN + 'All data is in ELS, it\'s ok' + bcolors.ENDC)
 
     return file_nb_line == els_nb_doc