From d9d643333a8be6a432a73a776f35d679a3f7457f Mon Sep 17 00:00:00 2001
From: Matthew Neeley <maffoo@google.com>
Date: Fri, 8 Apr 2022 18:32:53 -0700
Subject: [PATCH] Add q helper function for constructing common qubit types.
 (#5181)

* Add q helper function

* fmt

* Ignore static type errors in tests

* Ignore intentional type errors

* Fix Raises docstring

* Fixes from review

* Use *args

* remove unused import
---
 cirq-core/cirq/__init__.py          |  1 +
 cirq-core/cirq/ops/__init__.py      |  2 +
 cirq-core/cirq/ops/qid_util.py      | 68 +++++++++++++++++++++++++++++
 cirq-core/cirq/ops/qid_util_test.py | 31 +++++++++++++
 4 files changed, 102 insertions(+)
 create mode 100644 cirq-core/cirq/ops/qid_util.py
 create mode 100644 cirq-core/cirq/ops/qid_util_test.py

diff --git a/cirq-core/cirq/__init__.py b/cirq-core/cirq/__init__.py
index 2410332da08..19c9c3009be 100644
--- a/cirq-core/cirq/__init__.py
+++ b/cirq-core/cirq/__init__.py
@@ -290,6 +290,7 @@
     ProjectorString,
     ProjectorSum,
     RandomGateChannel,
+    q,
     qft,
     Qid,
     QuantumFourierTransformGate,
diff --git a/cirq-core/cirq/ops/__init__.py b/cirq-core/cirq/ops/__init__.py
index 94e205fe185..b882c4be733 100644
--- a/cirq-core/cirq/ops/__init__.py
+++ b/cirq-core/cirq/ops/__init__.py
@@ -258,6 +258,8 @@
     PhasedXZGate,
 )
 
+from cirq.ops.qid_util import q
+
 from cirq.ops.random_gate_channel import (
     RandomGateChannel,
 )
diff --git a/cirq-core/cirq/ops/qid_util.py b/cirq-core/cirq/ops/qid_util.py
new file mode 100644
index 00000000000..eb024ac3836
--- /dev/null
+++ b/cirq-core/cirq/ops/qid_util.py
@@ -0,0 +1,68 @@
+# Copyright 2022 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.
+
+from typing import overload, TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    import cirq
+
+
+@overload
+def q(__x: int) -> 'cirq.LineQubit':
+    ...
+
+
+@overload
+def q(__row: int, __col: int) -> 'cirq.GridQubit':
+    ...
+
+
+@overload
+def q(__name: str) -> 'cirq.NamedQubit':
+    ...
+
+
+def q(*args: Union[int, str]) -> Union['cirq.LineQubit', 'cirq.GridQubit', 'cirq.NamedQubit']:
+    """Constructs a qubit id of the appropriate type based on args.
+
+    This is shorthand for constructing qubit ids of common types:
+    >>> cirq.q(1) == cirq.LineQubit(1)
+    >>> cirq.q(1, 2) == cirq.GridQubit(1, 2)
+    >>> cirq.q("foo") == cirq.NamedQubit("foo")
+
+    Note that arguments should be treated as positional only, even
+    though this is only enforceable in python 3.8 or later.
+
+    Args:
+        *args: One or two ints, or a single str, as described above.
+
+    Returns:
+        cirq.LineQubit if called with one integer arg.
+        cirq.GridQubit if called with two integer args.
+        cirq.NamedQubit if called with one string arg.
+
+    Raises:
+        ValueError: if called with invalid arguments.
+    """
+    import cirq  # avoid circular import
+
+    if len(args) == 1:
+        if isinstance(args[0], int):
+            return cirq.LineQubit(args[0])
+        elif isinstance(args[0], str):
+            return cirq.NamedQubit(args[0])
+    elif len(args) == 2:
+        if isinstance(args[0], int) and isinstance(args[1], int):
+            return cirq.GridQubit(args[0], args[1])
+    raise ValueError(f"Could not construct qubit: args={args}")
diff --git a/cirq-core/cirq/ops/qid_util_test.py b/cirq-core/cirq/ops/qid_util_test.py
new file mode 100644
index 00000000000..446c6d59f4e
--- /dev/null
+++ b/cirq-core/cirq/ops/qid_util_test.py
@@ -0,0 +1,31 @@
+# Copyright 2022 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.
+
+import pytest
+
+import cirq
+
+
+def test_q() -> None:
+    assert cirq.q(0) == cirq.LineQubit(0)
+    assert cirq.q(1, 2) == cirq.GridQubit(1, 2)
+    assert cirq.q("foo") == cirq.NamedQubit("foo")
+
+
+def test_q_invalid() -> None:
+    # Ignore static type errors so we can test runtime typechecks.
+    with pytest.raises(ValueError):
+        cirq.q([1, 2, 3])  # type: ignore[call-overload]
+    with pytest.raises(ValueError):
+        cirq.q(1, "foo")  # type: ignore[call-overload]