Mypy: Annotate PauliTerm

CHANGELOG.md

@@ -29,7 +29,8 @@ Changelog
     and it is now a keyword-only argument (@karalekas, gh-1071).
 -   The code in `device.py` as been reorganized into a new `device` subdirectory
     in a completely backwards-compatible fashion (@karalekas, gh-1066).
-
+-   Type hints have been added to the `PauliTerm` class (@rht, gh-1075).
+
 ### Bugfixes
 
 -   Updated `examples/meyer_penny_game.py` with the correct path to the Meyer Penny

pyquil/paulis.py

@@ -22,7 +22,7 @@
 import numpy as np
 import copy
 
-from typing import Union
+from typing import Dict, List, Union
 
 from pyquil.quilatom import QubitPlaceholder
 
@@ -81,7 +81,7 @@ def __init__(self, op: str, index: int, coefficient: Union[int, float, complex]
         if op not in PAULI_OPS:
             raise ValueError(f"{op} is not a valid Pauli operator")
 
-        self._ops = OrderedDict()
+        self._ops: Dict[int, str] = OrderedDict()
         if op != "I":
             if not _valid_qubit(index):
                 raise ValueError(f"{index} is not a valid qubit")
@@ -90,7 +90,7 @@ def __init__(self, op: str, index: int, coefficient: Union[int, float, complex]
             raise ValueError("coefficient of PauliTerm must be a Number.")
         self.coefficient = complex(coefficient)
 
-    def id(self, sort_ops=True):
+    def id(self, sort_ops: bool = True) -> str:
         """
         Returns an identifier string for the PauliTerm (ignoring the coefficient).
 
@@ -130,7 +130,7 @@ def operations_as_set(self):
         """
         return frozenset(self._ops.items())
 
-    def __eq__(self, other):
+    def __eq__(self, other: object) -> bool:
         if not isinstance(other, (PauliTerm, PauliSum)):
             raise TypeError("Can't compare PauliTerm with object of type {}.".format(type(other)))
         elif isinstance(other, PauliSum):
@@ -146,14 +146,14 @@ def __hash__(self):
             self.operations_as_set()
         ))
 
-    def __len__(self):
+    def __len__(self) -> int:
         """
         The length of the PauliTerm is the number of Pauli operators in the term. A term that
         consists of only a scalar has a length of zero.
         """
         return len(self._ops)
 
-    def copy(self):
+    def copy(self) -> 'PauliTerm':
         """
         Properly creates a new PauliTerm, with a completely new dictionary
         of operators
@@ -173,19 +173,19 @@ def copy(self):
     def program(self):
         return Program([QUANTUM_GATES[gate](q) for q, gate in self])
 
-    def get_qubits(self):
+    def get_qubits(self) -> List[int]:
         """Gets all the qubits that this PauliTerm operates on.
         """
         return list(self._ops.keys())
 
-    def __getitem__(self, i):
+    def __getitem__(self, i: int) -> str:
         return self._ops.get(i, "I")
 
     def __iter__(self):
         for i in self.get_qubits():
             yield i, self[i]
 
-    def _multiply_factor(self, factor, index):
+    def _multiply_factor(self, factor: str, index: int) -> 'PauliTerm':
         new_term = PauliTerm("I", 0)
         new_coeff = self.coefficient
         new_ops = self._ops.copy()
@@ -203,13 +203,13 @@ def _multiply_factor(self, factor, index):
 
         return new_term
 
-    def __mul__(self, term):
+    def __mul__(self, term: Union['PauliTerm', 'PauliSum', 'Number']) -> Union['PauliTerm', 'PauliSum']:
         """Multiplies this Pauli Term with another PauliTerm, PauliSum, or number according to the
         Pauli algebra rules.
 
         :param term: (PauliTerm or PauliSum or Number) A term to multiply by.
         :returns: The product of this PauliTerm and term.
-        :rtype: PauliTerm
+        :rtype: PauliTerm or PauliSum
         """
         if isinstance(term, Number):
             return term_with_coeff(self, self.coefficient * term)
@@ -224,7 +224,7 @@ def __mul__(self, term):
 
             return term_with_coeff(new_term, new_term.coefficient * new_coeff)
 
-    def __rmul__(self, other):
+    def __rmul__(self, other: Number) -> 'PauliTerm':
         """Multiplies this PauliTerm with another object, probably a number.
 
         :param other: A number or PauliTerm to multiply by
@@ -234,7 +234,7 @@ def __rmul__(self, other):
         assert isinstance(other, Number)
         return self * other
 
-    def __pow__(self, power):
+    def __pow__(self, power: int) -> 'PauliTerm':
         """Raises this PauliTerm to power.
 
         :param int power: The power to raise this PauliTerm to.
@@ -253,10 +253,10 @@ def __pow__(self, power):
             result *= self
         return result
 
-    def __add__(self, other):
+    def __add__(self, other: Union['PauliTerm', 'PauliSum', Number]) -> 'PauliSum':
         """Adds this PauliTerm with another one.
 
-        :param other: A PauliTerm object or a Number
+        :param other: A PauliTerm object, a PauliSum object, or a Number
         :returns: A PauliSum object representing the sum of this PauliTerm and other
         :rtype: PauliSum
         """
@@ -268,7 +268,7 @@ def __add__(self, other):
             new_sum = PauliSum([self, other])
             return new_sum.simplify()
 
-    def __radd__(self, other):
+    def __radd__(self, other: Number) -> 'PauliTerm':
         """Adds this PauliTerm with a Number.
 
         :param other: A PauliTerm object or a Number
@@ -278,7 +278,7 @@ def __radd__(self, other):
         assert isinstance(other, Number)
         return PauliTerm("I", 0, other) + self
 
-    def __sub__(self, other):
+    def __sub__(self, other: Union['PauliTerm', Number]) -> 'PauliSum':
         """Subtracts a PauliTerm from this one.
 
         :param other: A PauliTerm object or a Number
@@ -287,7 +287,7 @@ def __sub__(self, other):
         """
         return self + -1. * other
 
-    def __rsub__(self, other):
+    def __rsub__(self, other: Union['PauliTerm', Number]) -> 'PauliSum':
         """Subtracts this PauliTerm from a Number or PauliTerm.
 
         :param other: A PauliTerm object or a Number
@@ -296,7 +296,7 @@ def __rsub__(self, other):
         """
         return other + -1. * self
 
-    def __str__(self):
+    def __str__(self) -> str:
         term_strs = []
         for index in self._ops.keys():
             term_strs.append("%s%s" % (self[index], index))
@@ -306,7 +306,7 @@ def __str__(self):
         out = "%s*%s" % (self.coefficient, '*'.join(term_strs))
         return out
 
-    def compact_str(self):
+    def compact_str(self) -> str:
         """A string representation of the Pauli term that is more compact than ``str(term)``
 
         >>> term = 2.0 * sX(1)* sZ(2)
@@ -318,7 +318,7 @@ def compact_str(self):
         return f'{self.coefficient}*{self.id(sort_ops=False)}'
 
     @classmethod
-    def from_list(cls, terms_list, coefficient=1.0):
+    def from_list(cls, terms_list, coefficient: float = 1.0) -> 'PauliTerm':
         """
         Allocates a Pauli Term from a list of operators and indices. This is more efficient than
         multiplying together individual terms.
@@ -351,7 +351,7 @@ def from_list(cls, terms_list, coefficient=1.0):
         return pterm
 
     @classmethod
-    def from_compact_str(cls, str_pauli_term):
+    def from_compact_str(cls, str_pauli_term: str) -> 'PauliTerm':
         """Construct a PauliTerm from the result of str(pauli_term)
         """
         # split into str_coef, str_op at first '*'' outside parenthesis
@@ -389,7 +389,7 @@ def from_compact_str(cls, str_pauli_term):
 
         return op
 
-    def pauli_string(self, qubits=None):
+    def pauli_string(self, qubits=None) -> str:
         """
         Return a string representation of this PauliTerm without its coefficient and with
         implicit qubit indices.
@@ -421,21 +421,21 @@ def pauli_string(self, qubits=None):
 
 
 # For convenience, a shorthand for several operators.
-def ID():
+def ID() -> PauliTerm:
     """
     The identity operator.
     """
     return PauliTerm("I", 0, 1)
 
 
-def ZERO():
+def ZERO() -> PauliTerm:
     """
     The zero operator.
     """
     return PauliTerm("I", 0, 0)
 
 
-def sI(q=None):
+def sI(q: int = None) -> PauliTerm:
     """
     A function that returns the identity operator, optionally on a particular qubit.
 
@@ -448,7 +448,7 @@ def sI(q=None):
     return PauliTerm("I", q)
 
 
-def sX(q):
+def sX(q: int) -> PauliTerm:
     """
     A function that returns the sigma_X operator on a particular qubit.
 
@@ -459,7 +459,7 @@ def sX(q):
     return PauliTerm("X", q)
 
 
-def sY(q):
+def sY(q: int) -> PauliTerm:
     """
     A function that returns the sigma_Y operator on a particular qubit.
 
@@ -470,7 +470,7 @@ def sY(q):
     return PauliTerm("Y", q)
 
 
-def sZ(q):
+def sZ(q: int) -> PauliTerm:
     """
     A function that returns the sigma_Z operator on a particular qubit.
 
@@ -481,7 +481,7 @@ def sZ(q):
     return PauliTerm("Z", q)
 
 
-def term_with_coeff(term, coeff):
+def term_with_coeff(term: PauliTerm, coeff: Number) -> PauliTerm:
     """
     Change the coefficient of a PauliTerm.