Skip to content

Commit 5d59cf7

Browse files
authored
Ensure InsertStrategy.NAME is identical after pickle round-trip (#7138)
The `Circuit.insert` uses identity to determine the `InsertStrategy` kind. Here we make InsertStrategy.NAME to get restored to the same object after pickle-unpickle round trip. This should prevent unexpected behavior if `cirq.Insert` is used in a multiprocessing call.
1 parent abf6051 commit 5d59cf7

File tree

2 files changed

+29
-0
lines changed

2 files changed

+29
-0
lines changed

cirq-core/cirq/circuits/insert_strategy.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,16 @@ class InsertStrategy:
2323
INLINE: 'InsertStrategy'
2424
EARLIEST: 'InsertStrategy'
2525

26+
def __new__(cls, name: str, doc: str) -> 'InsertStrategy':
27+
inst = getattr(cls, name, None)
28+
if not inst or not isinstance(inst, cls):
29+
inst = super().__new__(cls)
30+
return inst
31+
32+
def __getnewargs__(self):
33+
"""Returns a tuple of args to pass to __new__ when unpickling."""
34+
return (self.name, self.__doc__)
35+
2636
def __init__(self, name: str, doc: str):
2737
self.name = name
2838
self.__doc__ = doc

cirq-core/cirq/circuits/insert_strategy_test.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,27 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import pickle
16+
17+
import pytest
18+
1519
import cirq
1620

1721

1822
def test_repr():
1923
assert repr(cirq.InsertStrategy.NEW) == 'cirq.InsertStrategy.NEW'
24+
25+
26+
@pytest.mark.parametrize(
27+
'strategy',
28+
[
29+
cirq.InsertStrategy.NEW,
30+
cirq.InsertStrategy.NEW_THEN_INLINE,
31+
cirq.InsertStrategy.INLINE,
32+
cirq.InsertStrategy.EARLIEST,
33+
],
34+
ids=lambda strategy: strategy.name,
35+
)
36+
def test_identity_after_pickling(strategy: cirq.InsertStrategy):
37+
unpickled_strategy = pickle.loads(pickle.dumps(strategy))
38+
assert unpickled_strategy is strategy

0 commit comments

Comments
 (0)