Fix mypy and arrange utils

microsoft · May 20, 2022 · 883f2b5 · 883f2b5
1 parent 53b41a1
commit 883f2b5
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 10 deletions.
diff --git a/quantum-viz/pyproject.toml b/quantum-viz/pyproject.toml
@@ -73,7 +73,7 @@ show_error_codes = true
 show_error_context = true
 
 [[tool.mypy.overrides]]
-module = ["IPython.core.display", "varname.*"]
+module = ["IPython.*", "varname.*", "qiskit.*"]
 ignore_missing_imports = true
 
 [build-system]

diff --git a/quantum-viz/quantum_viz/qiskit_parser.py b/quantum-viz/quantum_viz/qiskit_parser.py
@@ -6,6 +6,7 @@
 from typing import Iterator
 from typing import List
 from typing import Optional
+from typing import Tuple
 from typing import Type
 from typing import Union
 
@@ -133,7 +134,7 @@ def __init__(
         }
         self.qubit2id: Dict[Qubit, int] = dict()
         self._init_qubits()
-        self._clbit2id: Dict[Clbit, (int, int)] = dict()
+        self._clbit2id: Dict[Clbit, Tuple[int, int]] = dict()
         self._update_qviz_dict()
 
     def _init_qubits(self) -> None:
@@ -210,7 +211,7 @@ def _parse_operation(
         depth: int,
     ) -> Dict:
 
-        op_dict = {"gate": self._get_instruction_name(instruction)}
+        op_dict: Dict = {"gate": self._get_instruction_name(instruction)}
 
         if instruction.params:
             self._add_params(op_dict, instruction)
@@ -348,7 +349,7 @@ def _update_condition(self, op_dict: Dict, instruction: Instruction) -> Dict:
         if isinstance(classical, ClassicalRegister):
             clbit = Clbit(classical, index=0)
         else:
-            clbit: Clbit = classical
+            clbit = classical
 
         if val:
             render_condition = ConditionalRender.ON_ONE.value

diff --git a/quantum-viz/quantum_viz/utils.py b/quantum-viz/quantum_viz/utils.py
@@ -11,7 +11,8 @@
 from typing import TYPE_CHECKING
 from typing import Union
 
-from widget import Viewer
+from .widget import DEFAULT_STYLE
+from .widget import Style
 
 if TYPE_CHECKING:
     from qiskit import QuantumCircuit
@@ -44,13 +45,18 @@
 </html>
 """
 SUFFIX = "_qviz.html"
-STYLES = ("Default", "BlackAndWhite", "Inverted")  # Use Literal type in python 3.8
+
+STYLES = list(Style)
+
+
+class UnsupportedStyleWarning(UserWarning):
+    """A warning raised when an unsupported style is chosen."""
 
 
 def display(
     circuit: Union[Dict[str, Any], "QuantumCircuit"],
     filename: Union[str, Path, None] = None,
-    style: str = "Default",
+    style: Style = DEFAULT_STYLE,
     version: Optional[str] = None,
     **kwargs,
 ) -> None:
@@ -65,7 +71,8 @@ def display(
     if style not in STYLES:
         warnings.warn(
             f"The selected style '{style}' is not supported and will be ignored.\n"
-            f"The supported styles are {STYLES}"
+            f"The supported styles are {STYLES}",
+            UnsupportedStyleWarning,
         )
 
     if version is None:
@@ -74,9 +81,9 @@ def display(
         version = "@" + version
 
     if not isinstance(circuit, dict):
-        from quantum_viz.qiskit_parser import qiskit2dict
+        from .qiskit_parser import qiskit2dict
 
-        circuit = qiskit2dict(circuit)
+        circuit = qiskit2dict(circuit, **kwargs)
 
     qviz_json = json.dumps(circuit)