From e40dc3ce2c2260fe6c39aa5d31dc98305580e8e4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= <benoitsierro@bluewin.ch>
Date: Mon, 31 May 2021 14:56:00 +0200
Subject: [PATCH] hush

---
 src/scgenerator/const.py                     |  1 +
 src/scgenerator/data/submit_job_template.txt |  1 +
 src/scgenerator/physics/simulate.py          |  3 +--
 src/scgenerator/utils.py                     | 14 +++++++++-----
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/scgenerator/const.py b/src/scgenerator/const.py
index 4496819..0dcf2f8 100644
--- a/src/scgenerator/const.py
+++ b/src/scgenerator/const.py
@@ -246,6 +246,7 @@ valid_variable = dict(
 )
 
 ENVIRON_KEY_BASE = "SCGENERATOR_"
+HUSH_PROGRESS = ENVIRON_KEY_BASE + "HUSH_PROGRESS"
 TMP_FOLDER_KEY_BASE = ENVIRON_KEY_BASE + "SC_TMP_"
 PREFIX_KEY_BASE = ENVIRON_KEY_BASE + "PREFIX_"
 PARAM_SEPARATOR = " "
diff --git a/src/scgenerator/data/submit_job_template.txt b/src/scgenerator/data/submit_job_template.txt
index 609b86f..fba5c55 100644
--- a/src/scgenerator/data/submit_job_template.txt
+++ b/src/scgenerator/data/submit_job_template.txt
@@ -50,5 +50,6 @@ done
 ##############################################################################################
 
 #### call your code below
+export SCGENERATOR_HUSH_PROGRESS=""
 scgenerator {command} {configs_list}
 exit
\ No newline at end of file
diff --git a/src/scgenerator/physics/simulate.py b/src/scgenerator/physics/simulate.py
index 232d4c7..1cf4e5f 100644
--- a/src/scgenerator/physics/simulate.py
+++ b/src/scgenerator/physics/simulate.py
@@ -2,10 +2,9 @@ import multiprocessing
 import os
 import sys
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type
+from typing import Any, Dict, List, Tuple, Type
 
 import numpy as np
-from numba import jit
 from tqdm import tqdm
 
 from .. import initialize, io, utils, const
diff --git a/src/scgenerator/utils.py b/src/scgenerator/utils.py
index 6c10314..bcf9901 100644
--- a/src/scgenerator/utils.py
+++ b/src/scgenerator/utils.py
@@ -10,10 +10,11 @@ import datetime as dt
 import itertools
 import logging
 import multiprocessing
-import re
 import socket
-from typing import Any, Callable, Dict, Iterator, List, Mapping, Tuple, Union
+import os
+from typing import Any, Dict, Iterator, List, Mapping, Tuple, Union
 from asyncio import Event
+from io import StringIO
 
 import numpy as np
 import ray
@@ -21,7 +22,7 @@ from copy import deepcopy
 
 from tqdm import tqdm
 
-from .const import PARAM_SEPARATOR, PREFIX_KEY_BASE, valid_variable, pbar_format
+from .const import PARAM_SEPARATOR, PREFIX_KEY_BASE, valid_variable, pbar_format, HUSH_PROGRESS
 from .logger import get_logger
 from .math import *
 
@@ -175,8 +176,11 @@ def progress_worker(num_steps: int, progress_queue: multiprocessing.Queue):
             Literal[0] : stop the worker and close the progress bars
             Tuple[int, float] : worker id and relative progress between 0 and 1
     """
+    kwargs = {}
+    if os.getenv(HUSH_PROGRESS) is not None:
+        kwargs = dict(file=StringIO())
     pbars: Dict[int, tqdm] = {}
-    with tqdm(total=num_steps, desc="Simulating", unit="step", position=0) as tq:
+    with tqdm(total=num_steps, desc="Simulating", unit="step", position=0, **kwargs) as tq:
         while True:
             raw = progress_queue.get()
             if raw == 0:
@@ -185,7 +189,7 @@ def progress_worker(num_steps: int, progress_queue: multiprocessing.Queue):
                 return
             i, rel_pos = raw
             if i not in pbars:
-                pbars[i] = tqdm(**pbar_format(i))
+                pbars[i] = tqdm(**pbar_format(i), **kwargs)
             pbars[i].update(rel_pos - pbars[i].n)
             tq.update()