Extend include_tags to support by-type selection (#7425)

implements #2905 based on the comments
[here](#2905 (comment)).
closed previous PR #7415

Author: WingCode

cirq-core/cirq/circuits/circuit.py

@@ -1172,7 +1172,7 @@ def to_text_diagram(
         *,
         use_unicode_characters: bool = True,
         transpose: bool = False,
-        include_tags: bool = True,
+        include_tags: bool | Iterable[type] = True,
         precision: int | None = 3,
         qubit_order: cirq.QubitOrderOrList = ops.QubitOrder.DEFAULT,
     ) -> str:
@@ -1182,7 +1182,10 @@ def to_text_diagram(
             use_unicode_characters: Determines if unicode characters are
                 allowed (as opposed to ascii-only diagrams).
             transpose: Arranges qubit wires vertically instead of horizontally.
-            include_tags: Whether tags on TaggedOperations should be printed
+            include_tags: Controls which tags attached to operations are
+                included. ``True`` includes all tags, ``False`` includes none,
+                or a collection of tag classes may be specified to include only
+                those tags.
             precision: Number of digits to display in text diagram
             qubit_order: Determines how qubits are ordered in the diagram.
 
@@ -1209,7 +1212,7 @@ def to_text_diagram_drawer(
         use_unicode_characters: bool = True,
         qubit_namer: Callable[[cirq.Qid], str] | None = None,
         transpose: bool = False,
-        include_tags: bool = True,
+        include_tags: bool | Iterable[type] = True,
         draw_moment_groups: bool = True,
         precision: int | None = 3,
         qubit_order: cirq.QubitOrderOrList = ops.QubitOrder.DEFAULT,
@@ -1224,7 +1227,10 @@ def to_text_diagram_drawer(
                 allowed (as opposed to ascii-only diagrams).
             qubit_namer: Names qubits in diagram. Defaults to using _circuit_diagram_info_ or str.
             transpose: Arranges qubit wires vertically instead of horizontally.
-            include_tags: Whether to include tags in the operation.
+            include_tags: Controls which tags attached to operations are
+                included. ``True`` includes all tags, ``False`` includes none,
+                or a collection of tag classes may be specified to include only
+                those tags.
             draw_moment_groups: Whether to draw moment symbol or not
             precision: Number of digits to use when representing numbers.
             qubit_order: Determines how qubits are ordered in the diagram.
@@ -2534,7 +2540,7 @@ def _draw_moment_annotations(
     get_circuit_diagram_info: Callable[
         [cirq.Operation, cirq.CircuitDiagramInfoArgs], cirq.CircuitDiagramInfo
     ],
-    include_tags: bool,
+    include_tags: bool | Iterable[type],
     first_annotation_row: int,
     transpose: bool,
 ):
@@ -2566,7 +2572,7 @@ def _draw_moment_in_diagram(
     get_circuit_diagram_info: (
         Callable[[cirq.Operation, cirq.CircuitDiagramInfoArgs], cirq.CircuitDiagramInfo] | None
     ),
-    include_tags: bool,
+    include_tags: bool | Iterable[type],
     first_annotation_row: int,
     transpose: bool,
 ):
@@ -2637,8 +2643,16 @@ def _draw_moment_in_diagram(
         desc = _formatted_phase(global_phase, use_unicode_characters, precision)
         if desc:
             y = max(label_map.values(), default=0) + 1
-            if tags and include_tags:
-                desc = desc + f"[{', '.join(map(str, tags))}]"
+            visible_tags = protocols.CircuitDiagramInfoArgs(
+                known_qubits=None,
+                known_qubit_count=None,
+                use_unicode_characters=True,
+                precision=None,
+                label_map=None,
+                include_tags=include_tags,
+            ).tags_to_include(tags)
+            if visible_tags:
+                desc = desc + f"[{', '.join(map(str, visible_tags))}]"
             out_diagram.write(x0, y, desc)
 
     if not non_global_ops:

cirq-core/cirq/circuits/moment.py

@@ -565,7 +565,7 @@ def to_text_diagram(
         extra_qubits: Iterable[cirq.Qid] = (),
         use_unicode_characters: bool = True,
         precision: int | None = None,
-        include_tags: bool = True,
+        include_tags: bool | Iterable[type] = True,
     ) -> str:
         """Create a text diagram for the moment.
 
@@ -583,8 +583,10 @@ def to_text_diagram(
             precision: How precise numbers, such as angles, should be. Use None
                 for infinite precision, or an integer for a certain number of
                 digits of precision.
-            include_tags: Whether or not to include operation tags in the
-                diagram.
+            include_tags: Controls which tags attached to operations are
+                included. ``True`` includes all tags, ``False`` includes none,
+                or a collection of tag classes may be specified to include only
+                those tags.
 
         Returns:
             The text diagram rendered into text.

cirq-core/cirq/contrib/quantum_volume/quantum_volume.py

@@ -127,7 +127,7 @@ def sample_heavy_set(
     # Add measure gates to the end of (a copy of) the circuit. Ensure that those
     # gates measure those in the given mapping, preserving this order.
     qubits = circuit.all_qubits()
-    key = None
+    key: Callable[[cirq.Qid], cirq.Qid] | None = None
     if mapping:
         # Add any qubits that were not explicitly mapped, so they aren't lost in
         # the sorting.
@@ -137,7 +137,7 @@ def sample_heavy_set(
     # Don't do a single large measurement gate because then the key will be one
     # large string. Instead, do a bunch of single-qubit measurement gates so we
     # preserve the qubit keys.
-    sorted_qubits = sorted(qubits, key=key)  # type: ignore[arg-type]
+    sorted_qubits = sorted(qubits, key=key)
     circuit_copy = circuit + [cirq.measure(q) for q in sorted_qubits]
 
     # Run the sampler to compare each output against the Heavy Set.

cirq-core/cirq/ops/raw_types.py

@@ -913,11 +913,12 @@ def _resolve_parameters_(
 
     def _circuit_diagram_info_(self, args: cirq.CircuitDiagramInfoArgs) -> cirq.CircuitDiagramInfo:
         sub_op_info = protocols.circuit_diagram_info(self.sub_operation, args, NotImplemented)
-        # Add tag to wire symbol if it exists.
-        if sub_op_info is not NotImplemented and args.include_tags and sub_op_info.wire_symbols:
-            sub_op_info.wire_symbols = (
-                sub_op_info.wire_symbols[0] + f"[{', '.join(map(str, self._tags))}]",
-            ) + sub_op_info.wire_symbols[1:]
+        if sub_op_info is not NotImplemented and sub_op_info.wire_symbols:
+            visible_tags = args.tags_to_include(self._tags)
+            if visible_tags:
+                sub_op_info.wire_symbols = (
+                    sub_op_info.wire_symbols[0] + f"[{', '.join(map(str, visible_tags))}]",
+                ) + sub_op_info.wire_symbols[1:]
         return sub_op_info
 
     @cached_method

cirq-core/cirq/ops/raw_types_test.py

@@ -554,6 +554,8 @@ def __str__(self):
     diagram_with_non_string_tag = "(1, 1): ───H[<taggy>]───"
     assert c.to_text_diagram() == diagram_with_non_string_tag
     assert c.to_text_diagram(include_tags=False) == diagram_without_tags
+    assert c.to_text_diagram(include_tags={str}) == diagram_without_tags
+    assert c.to_text_diagram(include_tags={TaggyTag}) == diagram_with_non_string_tag
 
 
 def test_circuit_diagram_tagged_global_phase() -> None:

cirq-core/cirq/protocols/circuit_diagram_info_protocol.py

@@ -181,7 +181,10 @@ class CircuitDiagramInfoArgs:
         precision: The number of digits after the decimal to show for numbers in
             the text diagram. None means use full precision.
         label_map: The map from label entities to diagram positions.
-        include_tags: Whether to print tags from TaggedOperations.
+        include_tags: If ``True`` all tags from ``TaggedOperations`` will be
+            printed.  If ``False`` no tags will be printed.  Alternatively a
+            collection of tag classes can be provided.  In this case only tags
+            whose type is contained in the collection will be shown.
         transpose: Whether the circuit is to be drawn with time from left to
             right (transpose is False), or from top to bottom.
     """
@@ -195,15 +198,19 @@ def __init__(
         use_unicode_characters: bool,
         precision: int | None,
         label_map: dict[LabelEntity, int] | None,
-        include_tags: bool = True,
+        include_tags: bool | Iterable[type] = True,
         transpose: bool = False,
     ) -> None:
         self.known_qubits = None if known_qubits is None else tuple(known_qubits)
         self.known_qubit_count = known_qubit_count
         self.use_unicode_characters = use_unicode_characters
         self.precision = precision
         self.label_map = label_map
-        self.include_tags = include_tags
+        self.include_tags: bool | frozenset[type]
+        if isinstance(include_tags, bool):
+            self.include_tags = include_tags
+        else:
+            self.include_tags = frozenset(include_tags)
         self.transpose = transpose
 
     def _value_equality_values_(self) -> Any:
@@ -217,7 +224,11 @@ def _value_equality_values_(self) -> Any:
                 if self.label_map is None
                 else tuple(sorted(self.label_map.items(), key=lambda e: e[0]))
             ),
-            self.include_tags,
+            (
+                self.include_tags
+                if isinstance(self.include_tags, bool)
+                else tuple(sorted(self.include_tags, key=lambda c: c.__name__))
+            ),
             self.transpose,
         )
 
@@ -229,10 +240,27 @@ def __repr__(self) -> str:
             f'use_unicode_characters={self.use_unicode_characters!r}, '
             f'precision={self.precision!r}, '
             f'label_map={self.label_map!r}, '
-            f'include_tags={self.include_tags!r}, '
+            f'include_tags={self._include_tags_repr()}, '
             f'transpose={self.transpose!r})'
         )
 
+    def _include_tags_repr(self) -> str:
+        if isinstance(self.include_tags, bool):
+            return repr(self.include_tags)
+        items = []
+        for cls in self.include_tags:
+            if cls.__module__ == 'builtins':
+                items.append(cls.__qualname__)
+            else:
+                items.append(f"{cls.__module__}.{cls.__qualname__}")
+        joined = ', '.join(items)
+        return f'frozenset({{{joined}}})'
+
+    def tags_to_include(self, tags: Iterable[Any]) -> list[Any]:
+        if isinstance(self.include_tags, bool):
+            return list(tags) if self.include_tags else []
+        return [t for t in tags if any(isinstance(t, cls) for cls in self.include_tags)]
+
     def format_real(self, val: sympy.Basic | int | float) -> str:
         if isinstance(val, sympy.Basic):
             return str(val)
@@ -279,6 +307,7 @@ def copy(self):
             use_unicode_characters=self.use_unicode_characters,
             precision=self.precision,
             label_map=self.label_map,
+            include_tags=self.include_tags,
             transpose=self.transpose,
         )
 

cirq-core/cirq/protocols/circuit_diagram_info_protocol_test.py

@@ -21,6 +21,10 @@
 import cirq
 
 
+class CustomTag:
+    pass
+
+
 def test_circuit_diagram_info_value_wrapping() -> None:
     single_info = cirq.CircuitDiagramInfo(('Single',))
 
@@ -174,6 +178,26 @@ def test_circuit_diagram_info_args_eq() -> None:
             include_tags=False,
         )
     )
+    eq.add_equality_group(
+        cirq.CircuitDiagramInfoArgs(
+            known_qubits=cirq.LineQubit.range(2),
+            known_qubit_count=2,
+            use_unicode_characters=False,
+            precision=None,
+            label_map=None,
+            include_tags={str},
+        )
+    )
+    eq.add_equality_group(
+        cirq.CircuitDiagramInfoArgs(
+            known_qubits=cirq.LineQubit.range(2),
+            known_qubit_count=2,
+            use_unicode_characters=False,
+            precision=None,
+            label_map=None,
+            include_tags={CustomTag},
+        )
+    )
     eq.add_equality_group(
         cirq.CircuitDiagramInfoArgs(
             known_qubits=cirq.LineQubit.range(2),
@@ -208,6 +232,30 @@ def test_circuit_diagram_info_args_repr() -> None:
         )
     )
 
+    cirq.testing.assert_equivalent_repr(
+        cirq.CircuitDiagramInfoArgs(
+            known_qubits=cirq.LineQubit.range(1),
+            known_qubit_count=1,
+            use_unicode_characters=False,
+            precision=None,
+            label_map=None,
+            include_tags={str},
+            transpose=False,
+        )
+    )
+
+    cirq.testing.assert_equivalent_repr(
+        cirq.CircuitDiagramInfoArgs(
+            known_qubits=cirq.LineQubit.range(1),
+            known_qubit_count=1,
+            use_unicode_characters=False,
+            precision=None,
+            label_map=None,
+            include_tags={CustomTag},
+            transpose=False,
+        )
+    )
+
 
 def test_format_real() -> None:
     args = cirq.CircuitDiagramInfoArgs.UNINFORMED_DEFAULT.copy()