From 2eb8422d7842c56659ddee459dbc7c30305af431 Mon Sep 17 00:00:00 2001
From: rich <rich@super.tech>
Date: Mon, 23 Sep 2024 11:09:31 -0500
Subject: [PATCH] fix qcvv progress bars when using cli

---
 .../supermarq/qcvv/base_experiment.py              |  2 +-
 .../supermarq/qcvv/base_experiment_test.py         |  6 ------
 supermarq-benchmarks/supermarq/qcvv/conftest.py    | 14 ++++++++++++++
 supermarq-benchmarks/supermarq/qcvv/irb.py         |  4 ++--
 supermarq-benchmarks/supermarq/qcvv/irb_test.py    |  6 ------
 supermarq-benchmarks/supermarq/qcvv/xeb.py         |  5 ++---
 supermarq-benchmarks/supermarq/qcvv/xeb_test.py    |  6 ------
 7 files changed, 19 insertions(+), 24 deletions(-)
 create mode 100644 supermarq-benchmarks/supermarq/qcvv/conftest.py

diff --git a/supermarq-benchmarks/supermarq/qcvv/base_experiment.py b/supermarq-benchmarks/supermarq/qcvv/base_experiment.py
index a812ab25a..11dd9923e 100644
--- a/supermarq-benchmarks/supermarq/qcvv/base_experiment.py
+++ b/supermarq-benchmarks/supermarq/qcvv/base_experiment.py
@@ -24,7 +24,7 @@
 import cirq
 import cirq_superstaq as css
 import pandas as pd
-from tqdm.notebook import tqdm
+from tqdm.auto import tqdm
 
 
 @dataclass
diff --git a/supermarq-benchmarks/supermarq/qcvv/base_experiment_test.py b/supermarq-benchmarks/supermarq/qcvv/base_experiment_test.py
index ee7d2fac9..dc3095db2 100644
--- a/supermarq-benchmarks/supermarq/qcvv/base_experiment_test.py
+++ b/supermarq-benchmarks/supermarq/qcvv/base_experiment_test.py
@@ -17,7 +17,6 @@
 
 from __future__ import annotations
 
-import os
 from dataclasses import dataclass
 from unittest.mock import MagicMock, call, patch
 
@@ -38,11 +37,6 @@ class ExampleResults(BenchmarkingResults):
     experiment_name = "Example results"
 
 
-@pytest.fixture(scope="session", autouse=True)
-def patch_tqdm() -> None:
-    os.environ["TQDM_DISABLE"] = "1"
-
-
 @pytest.fixture
 @patch.multiple(BenchmarkingExperiment, __abstractmethods__=set())
 def abc_experiment() -> BenchmarkingExperiment[ExampleResults]:
diff --git a/supermarq-benchmarks/supermarq/qcvv/conftest.py b/supermarq-benchmarks/supermarq/qcvv/conftest.py
new file mode 100644
index 000000000..4bf01245c
--- /dev/null
+++ b/supermarq-benchmarks/supermarq/qcvv/conftest.py
@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+
+
+@pytest.fixture(scope="session", autouse=True)
+def patch_tqdm() -> Generator[None, None, None]:
+    """Disable progress bars during tests."""
+    with mock.patch.dict(os.environ, {"TQDM_DISABLE": "1"}):
+        yield
diff --git a/supermarq-benchmarks/supermarq/qcvv/irb.py b/supermarq-benchmarks/supermarq/qcvv/irb.py
index ce39957c7..553e9136f 100644
--- a/supermarq-benchmarks/supermarq/qcvv/irb.py
+++ b/supermarq-benchmarks/supermarq/qcvv/irb.py
@@ -26,8 +26,8 @@
 import pandas as pd
 import scipy
 import seaborn as sns
+from tqdm.auto import trange
 from tqdm.contrib.itertools import product
-from tqdm.notebook import tqdm
 
 from supermarq.qcvv.base_experiment import BenchmarkingExperiment, BenchmarkingResults, Sample
 
@@ -408,7 +408,7 @@ def gates_per_clifford(self, samples: int = 500) -> dict[str, float]:
         """
         sample = [
             self._clifford_gate_to_circuit(self.random_clifford())
-            for _ in tqdm(range(samples), desc="Sampling Clifford operations")
+            for _ in trange(samples, desc="Sampling Clifford operations")
         ]
         return {
             "single_qubit_gates": np.mean(
diff --git a/supermarq-benchmarks/supermarq/qcvv/irb_test.py b/supermarq-benchmarks/supermarq/qcvv/irb_test.py
index 37224c6f7..ae8e66e95 100644
--- a/supermarq-benchmarks/supermarq/qcvv/irb_test.py
+++ b/supermarq-benchmarks/supermarq/qcvv/irb_test.py
@@ -16,7 +16,6 @@
 # mypy: disable-error-code="union-attr"
 from __future__ import annotations
 
-import os
 from unittest.mock import MagicMock, patch
 
 import cirq
@@ -28,11 +27,6 @@
 from supermarq.qcvv.irb import IRB
 
 
-@pytest.fixture(scope="session", autouse=True)
-def patch_tqdm() -> None:
-    os.environ["TQDM_DISABLE"] = "1"
-
-
 def test_irb_init() -> None:
     with patch("cirq_superstaq.service.Service"):
         experiment = IRB()
diff --git a/supermarq-benchmarks/supermarq/qcvv/xeb.py b/supermarq-benchmarks/supermarq/qcvv/xeb.py
index 74f3f0372..df8b0a4fc 100644
--- a/supermarq-benchmarks/supermarq/qcvv/xeb.py
+++ b/supermarq-benchmarks/supermarq/qcvv/xeb.py
@@ -24,8 +24,7 @@
 import pandas as pd
 import scipy
 import seaborn as sns
-import tqdm
-import tqdm.contrib
+import tqdm.auto
 import tqdm.contrib.itertools
 
 from supermarq.qcvv import BenchmarkingExperiment, BenchmarkingResults, Sample
@@ -252,7 +251,7 @@ def _process_probabilities(
             sample.sample_probabilities = sample.probabilities
             sample.probabilities = {}
 
-        for sample in tqdm.notebook.tqdm(samples, desc="Evaluating circuits"):
+        for sample in tqdm.auto.tqdm(samples, desc="Evaluating circuits"):
             sample.target_probabilities = self._simulate_sample(sample)
 
         records = []
diff --git a/supermarq-benchmarks/supermarq/qcvv/xeb_test.py b/supermarq-benchmarks/supermarq/qcvv/xeb_test.py
index 177d45aa6..61e035aaf 100644
--- a/supermarq-benchmarks/supermarq/qcvv/xeb_test.py
+++ b/supermarq-benchmarks/supermarq/qcvv/xeb_test.py
@@ -16,7 +16,6 @@
 from __future__ import annotations
 
 import itertools
-import os
 from unittest.mock import MagicMock, patch
 
 import cirq
@@ -27,11 +26,6 @@
 from supermarq.qcvv import XEB, XEBSample
 
 
-@pytest.fixture(scope="session", autouse=True)
-def patch_tqdm() -> None:
-    os.environ["TQDM_DISABLE"] = "1"
-
-
 def test_xeb_init() -> None:
     with patch("cirq_superstaq.service.Service"):
         experiment = XEB()