From 40a04d66bdd0dcb2ecfd9260625c4addf4f72475 Mon Sep 17 00:00:00 2001
From: Anurudh Peduri <anurudhp97@gmail.com>
Date: Tue, 20 Aug 2024 14:32:36 -0700
Subject: [PATCH] remove `_get_all_rotation_types`

---
 .../resource_counting/t_counts_from_sigma.py  | 30 ++-----------------
 .../t_counts_from_sigma_test.py               | 19 +-----------
 2 files changed, 4 insertions(+), 45 deletions(-)

diff --git a/qualtran/resource_counting/t_counts_from_sigma.py b/qualtran/resource_counting/t_counts_from_sigma.py
index 0d969544a5..b7e9ac39ab 100644
--- a/qualtran/resource_counting/t_counts_from_sigma.py
+++ b/qualtran/resource_counting/t_counts_from_sigma.py
@@ -23,38 +23,14 @@
     from qualtran import Bloq
 
 
-@runtime_checkable
-class _HasEps(Protocol):
-    """Protocol for typing `RotationBloq` base class mixin that has accuracy specified as eps."""
-
-    eps: float
-
-
-def _get_all_rotation_types() -> Tuple[Type['_HasEps'], ...]:
-    """Returns all classes defined in bloqs.basic_gates which have an attribute `eps`."""
-    from qualtran.bloqs.basic_gates import GlobalPhase
-
-    bloqs_to_exclude = [GlobalPhase]
-
-    return tuple(
-        cast(Type['_HasEps'], v)  # Can't use `issubclass` with protocols with attributes.
-        for (_, v) in inspect.getmembers(sys.modules['qualtran.bloqs.basic_gates'], inspect.isclass)
-        if isinstance(v, _HasEps) and v not in bloqs_to_exclude
-    )
-
-
-def t_counts_from_sigma(
-    sigma: Mapping['Bloq', SymbolicInt],
-    rotation_types: Optional[Tuple[Type['_HasEps'], ...]] = None,
-) -> SymbolicInt:
+def t_counts_from_sigma(sigma: Mapping['Bloq', SymbolicInt]) -> SymbolicInt:
     """Aggregates T-counts from a sigma dictionary by summing T-costs for all rotation bloqs."""
     from qualtran.bloqs.basic_gates import TGate
     from qualtran.cirq_interop.t_complexity_protocol import TComplexity
+    from qualtran.resource_counting.classify_bloqs import bloq_is_rotation
 
-    if rotation_types is None:
-        rotation_types = _get_all_rotation_types()
     ret = sigma.get(TGate(), 0) + sigma.get(TGate().adjoint(), 0)
     for bloq, counts in sigma.items():
-        if isinstance(bloq, rotation_types) and not cirq.has_stabilizer_effect(bloq):
+        if bloq_is_rotation(bloq) and not cirq.has_stabilizer_effect(bloq):
             ret += ceil(TComplexity.rotation_cost(bloq.eps)) * counts
     return ret
diff --git a/qualtran/resource_counting/t_counts_from_sigma_test.py b/qualtran/resource_counting/t_counts_from_sigma_test.py
index 181002c9e4..e374e408cc 100644
--- a/qualtran/resource_counting/t_counts_from_sigma_test.py
+++ b/qualtran/resource_counting/t_counts_from_sigma_test.py
@@ -19,7 +19,6 @@
     Rx,
     Ry,
     Rz,
-    SU2RotationGate,
     TGate,
     Toffoli,
     XPowGate,
@@ -27,23 +26,7 @@
     ZPowGate,
 )
 from qualtran.cirq_interop.t_complexity_protocol import TComplexity
-from qualtran.resource_counting.t_counts_from_sigma import (
-    _get_all_rotation_types,
-    t_counts_from_sigma,
-)
-
-
-def test_all_rotation_types():
-    assert set(_get_all_rotation_types()) == {
-        CZPowGate,
-        Rx,
-        Ry,
-        Rz,
-        XPowGate,
-        YPowGate,
-        ZPowGate,
-        SU2RotationGate,
-    }
+from qualtran.resource_counting.t_counts_from_sigma import t_counts_from_sigma
 
 
 def test_t_counts_from_sigma():