Source code for cirq.ops.op_tree

# Copyright 2018 The Cirq Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""A recursive type describing trees of operations, and utility methods for it.
"""

from typing import Callable, Iterable, Iterator, NoReturn, Union
from typing_extensions import Protocol

from cirq._doc import document
from cirq.ops.moment import Moment
from cirq.ops.raw_types import Operation


class OpTree(Protocol):
    """The recursive type consumed by circuit builder methods.

    An OpTree is a type protocol, satisfied by anything that can be recursively
    flattened into Operations. We also define the Union type OP_TREE which
    can be an OpTree or just a single Operation.

    For example:
    - An Operation is an OP_TREE all by itself.
    - A list of operations is an OP_TREE.
    - A list of tuples of operations is an OP_TREE.
    - A list with a mix of operations and lists of operations is an OP_TREE.
    - A generator yielding operations is an OP_TREE.

    Note: once mypy supports recursive types this could be defined as an alias:

    OP_TREE = Union[Operation, Iterable['OP_TREE']]

    See: https://github.com/python/mypy/issues/731
    """

    def __iter__(self) -> Iterator[Union[Operation, 'OpTree']]:
        pass


OP_TREE = Union[Operation, OpTree]
document(
    OP_TREE,  # type: ignore
    """An operation or nested collections of operations.

    Here are some examples of things that can be given to a method that takes a
    `cirq.OP_TREE` argument:

    - A single operation (a `cirq.Operation`).
    - A list of operations (a `List[cirq.Operation]`).
    - A list of lists of operations (a `List[List[cirq.Operation]]`).
    - A list mixing operations and generators of operations
        (a `List[Union[cirq.Operation, Iterator[cirq.Operation]]]`).
    - Generally anything that can be iterated, and its items iterated, and
        so forth recursively until a bottom layer of operations is found.
    """)


[docs]def flatten_op_tree(root: OP_TREE, preserve_moments: bool = False ) -> Iterator[Union[Operation, Moment]]: """Performs an in-order iteration of the operations (leaves) in an OP_TREE. Args: root: The operation or tree of operations to iterate. preserve_moments: Whether to yield Moments intact instead of flattening them Yields: Operations from the tree. Raises: TypeError: root isn't a valid OP_TREE. """ if preserve_moments: return flatten_to_ops_or_moments(root) else: return flatten_to_ops(root)
[docs]def flatten_to_ops(root: OP_TREE) -> Iterator[Operation]: """Performs an in-order iteration of the operations (leaves) in an OP_TREE. Args: root: The operation or tree of operations to iterate. Yields: Operations or moments from the tree. Raises: TypeError: root isn't a valid OP_TREE. """ if isinstance(root, Operation): yield root elif isinstance(root, Iterable) and not isinstance(root, str): for subtree in root: yield from flatten_to_ops(subtree) else: _bad_op_tree(root)
[docs]def flatten_to_ops_or_moments(root: OP_TREE ) -> Iterator[Union[Operation, Moment]]: """Performs an in-order iteration OP_TREE, yielding ops and moments. Args: root: The operation or tree of operations to iterate. Yields: Operations or moments from the tree. Raises: TypeError: root isn't a valid OP_TREE. """ if isinstance(root, (Operation, Moment)): yield root elif isinstance(root, Iterable) and not isinstance(root, str): for subtree in root: yield from flatten_to_ops_or_moments(subtree) else: _bad_op_tree(root)
[docs]def transform_op_tree( root: OP_TREE, op_transformation: Callable[[Operation], OP_TREE]=lambda e: e, iter_transformation: Callable[[Iterable[OP_TREE]], OP_TREE]=lambda e: e, preserve_moments: bool = False ) -> OP_TREE: """Maps transformation functions onto the nodes of an OP_TREE. Args: root: The operation or tree of operations to transform. op_transformation: How to transform the operations (i.e. leaves). iter_transformation: How to transform the iterables (i.e. internal nodes). preserve_moments: Whether to leave Moments alone. If True, the transformation functions will not be applied to Moments or the operations within them. Returns: A transformed operation tree. Raises: TypeError: root isn't a valid OP_TREE. """ if isinstance(root, Operation): return op_transformation(root) if preserve_moments and isinstance(root, Moment): return root if isinstance(root, Iterable) and not isinstance(root, str): return iter_transformation( transform_op_tree(subtree, op_transformation, iter_transformation, preserve_moments) for subtree in root) _bad_op_tree(root)
[docs]def freeze_op_tree(root: OP_TREE) -> OP_TREE: """Replaces all iterables in the OP_TREE with tuples. Args: root: The operation or tree of operations to freeze. Returns: An OP_TREE with the same operations and branching structure, but where all internal nodes are tuples instead of arbitrary iterables. """ return transform_op_tree(root, iter_transformation=tuple)
def _bad_op_tree(root: OP_TREE) -> NoReturn: raise TypeError(f'Not an Operation or Iterable: {type(root)} {root}')