From 352b3f8bf02e2fffeee525ba84f3e27fbdb5097a Mon Sep 17 00:00:00 2001
From: "Adam M. Krajewski" <54290107+amkrajewski@users.noreply.github.com>
Date: Sat, 24 Feb 2024 15:35:49 +0000
Subject: [PATCH] (core) imporved `Calculator` initialization printouts by
 conditional coloring of the output

---
 pyproject.toml              |  3 ++-
 pysipfenn/core/pysipfenn.py | 26 ++++++++++++++++----------
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 56940a7..006c1e9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,8 @@ dependencies = [
     "pymongo>=4.4",
     "pysmartdl2>=2.0.0",
     "dnspython",
-    "ruamel.yaml"
+    "ruamel.yaml",
+    "colorama"
 ]
 
 [project.optional-dependencies]
diff --git a/pysipfenn/core/pysipfenn.py b/pysipfenn/core/pysipfenn.py
index f58d04b..816fa31 100644
--- a/pysipfenn/core/pysipfenn.py
+++ b/pysipfenn/core/pysipfenn.py
@@ -12,6 +12,7 @@
 from tqdm.contrib.concurrent import process_map
 import natsort
 from pysmartdl2 import SmartDL
+from colorama import Fore, Style
 
 # Scientific Computing Imports
 import numpy as np
@@ -71,17 +72,17 @@ def __init__(self,
                  verbose: bool = True):
         """Initializes the pySIPFENN Calculator object."""
         if verbose:
-            print('*********  Initializing pySIPFENN Calculator  **********')
-        self.verbose = verbose
+            print('\n*********  Initializing pySIPFENN Calculator  **********')
+            self.verbose = verbose
         # dictionary with all model information
         with resources.files('pysipfenn.modelsSIPFENN').joinpath('models.json').open('r') as f:
             if verbose:
-                print(f'Loading model definitions from: {f.name}')
+                print(f'Loading model definitions from: {Fore.BLUE}{f.name}{Style.RESET_ALL}')
             self.models = json.load(f)
         # networks list
         self.network_list = list(self.models.keys())
         if verbose:
-            print(f'Found {len(self.network_list)} network definitions in models.json')
+            print(f'Found {Fore.BLUE}{len(self.network_list)} network definitions in models.json{Style.RESET_ALL}')
         # network names
         self.network_list_names = [self.models[net]['name'] for net in self.network_list]
         self.network_list_available = []
@@ -89,10 +90,10 @@ def __init__(self,
 
         self.loadedModels = {}
         if autoLoad:
-            print(f'Loading all available models (autoLoad=True)')
+            print(f'Loading all available models ({Fore.BLUE}autoLoad=True{Style.RESET_ALL})')
             self.loadModels()
         else:
-            print(f'Skipping model loading (autoLoad=False)')
+            print(f'Skipping model loading ({Fore.BLUE}autoLoad=False{Style.RESET_ALL})')
 
         self.prototypeLibrary = {}
         self.parsePrototypeLibrary(verbose=verbose)
@@ -105,7 +106,7 @@ def __init__(self,
         }
         self.inputFiles = []
         if verbose:
-            print(f'*********  pySIPFENN Successfully Initialized  **********')
+            print(f'{Fore.GREEN}**********      Successfully Initialized      **********{Style.RESET_ALL}')
 
     def __str__(self):
         """Prints the status of the ``Calculator`` object."""
@@ -175,7 +176,12 @@ class automatically initializes with ``verbose=True``.
                 }
             })
         if verbose:
-            print(f'{len(self.prototypeLibrary)} prototype structures present into the prototype library.')
+            protoLen = len(self.prototypeLibrary)
+            if protoLen == 0:
+                print(f"{Style.DIM}No prototypes were loaded into the prototype library.{Style.RESET_ALL}")
+            else:
+                print(f"Loaded {Fore.GREEN}{protoLen} prototypes {Style.RESET_ALL}into the library.")
+            
 
     def appendPrototypeLibrary(self, customPath: str) -> None:
         """Parses a custom prototype library YAML file and permanently appends it into the internal prototypeLibrary
@@ -205,9 +211,9 @@ def updateModelAvailability(self) -> None:
         for net, netName in zip(self.network_list, self.network_list_names):
             if all_files.__contains__(net + '.onnx'):
                 detectedNets.append(net)
-                print('✔ ' + netName)
+                print(f"{Fore.GREEN}✔ {netName}{Style.RESET_ALL}")
             else:
-                print('x ' + netName)
+                print(f"{Style.DIM}x {netName}{Style.RESET_ALL}")
         self.network_list_available = detectedNets
 
     def downloadModels(self, network: str = 'all') -> None: