1# Copyright 2018 The Cirq Developers
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import pytest
16
17import cirq
18
19
20def test_phase_by():
21    class NoMethod:
22        pass
23
24    class ReturnsNotImplemented:
25        def _phase_by_(self, phase_turns, qubit_on):
26            return NotImplemented
27
28    class PhaseIsAddition:
29        def __init__(self, num_qubits):
30            self.phase = [0] * num_qubits
31            self.num_qubits = num_qubits
32
33        def _phase_by_(self, phase_turns, qubit_on):
34            if qubit_on >= self.num_qubits:
35                return self
36            self.phase[qubit_on] += phase_turns
37            return self
38
39    n = NoMethod()
40    rin = ReturnsNotImplemented()
41
42    # Without default
43
44    with pytest.raises(TypeError, match='no _phase_by_ method'):
45        _ = cirq.phase_by(n, 1, 0)
46    with pytest.raises(TypeError, match='returned NotImplemented'):
47        _ = cirq.phase_by(rin, 1, 0)
48
49    # With default
50    assert cirq.phase_by(n, 1, 0, default=None) == None
51    assert cirq.phase_by(rin, 1, 0, default=None) == None
52
53    test = PhaseIsAddition(3)
54    assert test.phase == [0, 0, 0]
55    test = cirq.phase_by(test, 0.25, 0)
56    assert test.phase == [0.25, 0, 0]
57    test = cirq.phase_by(test, 0.25, 0)
58    assert test.phase == [0.50, 0, 0]
59    test = cirq.phase_by(test, 0.40, 1)
60    assert test.phase == [0.50, 0.40, 0]
61    test = cirq.phase_by(test, 0.40, 4)
62    assert test.phase == [0.50, 0.40, 0]
63