Repository: rigetticomputing/pyquil (python)

Issue:
get_qubits doesn't return qubits in Instr's
I would expect the following to include `0`.
```
>>> from pyquil.quil import Program
>>> p = Program()
>>> p.inst(("H",0))
<pyquil.quil.Program object at 0x7fc71e298290>
>>> p.get_qubits()
set([])
>>> 
```

Pull request:
Including Instr's in get_qubits
Fixes #155 

Patch:
diff --git a/pyquil/quil.py b/pyquil/quil.py
index bc94cf9..1939d67 100644
--- a/pyquil/quil.py
+++ b/pyquil/quil.py
@@ -16,14 +16,18 @@
 """
 Module for creating and defining Quil programs.
 """
+from six import integer_types
 
 from .quilbase import (InstructionGroup,
+                       Instr,
                        Addr,
                        While,
                        If,
                        DefGate,
                        Gate,
                        Measurement,
+                       issubinstance,
+                       AbstractQubit,
                        merge_resource_managers)
 
 from .gates import MEASURE, STANDARD_GATES
@@ -75,21 +79,38 @@ class Program(InstructionGroup):
 
     def get_qubits(self):
         """
-        :return: a set of all the qubit indices allocated in this program, synthesizing freely
-        allocated qubits if neccessary.
+        Returns all of the qubit indices used in this program, including gate applications and
+        allocated qubits. e.g.
+            >>> p = Program()
+            >>> p.inst(("H", 1))
+            >>> p.get_qubits()
+            {1}
+            >>> q = p.alloc()
+            >>> len(p.get_qubits())
+            2
+
+        :return: A set of all the qubit indices used in this program, synthesizing freely
+         allocated qubits if neccessary.
         :rtype: set
         """
-        allocated_qubits = set()
+        qubits = set()
         self.synthesize()
         for ii, action in self:
             if isinstance(action, Gate):
                 qubit_indices = {qq.index() for qq in action.arguments}
             elif isinstance(action, Measurement):
                 qubit_indices = {action.arguments[0].index()}
+            elif isinstance(action, Instr):
+                qubit_indices = set()
+                for arg in action.arguments:
+                    if isinstance(arg, integer_types):
+                        qubit_indices.add(arg)
+                    elif isinstance(arg, AbstractQubit):
+                        qubit_indices.add(arg.index())
             else:
                 continue
-            allocated_qubits = set.union(allocated_qubits, qubit_indices)
-        return allocated_qubits
+            qubits = set.union(qubits, qubit_indices)
+        return qubits
 
     def defgate(self, name, matrix):
         """
diff --git a/pyquil/quilbase.py b/pyquil/quilbase.py
index 26b09fa..239b965 100644
--- a/pyquil/quilbase.py
+++ b/pyquil/quilbase.py
@@ -668,6 +668,8 @@ class Instr(AbstractInstruction):
         for arg in self.arguments:
             if issubinstance(arg, AbstractQubit):
                 qubits.add(arg.index())
+            elif isinstance(arg, int):
+                qubits.add(arg)
         return qubits
 
 


Test patch:
diff --git a/pyquil/tests/test_quil.py b/pyquil/tests/test_quil.py
index 2383453..e328faf 100644
--- a/pyquil/tests/test_quil.py
+++ b/pyquil/tests/test_quil.py
@@ -462,3 +462,11 @@ def test_get_qubits():
     qq = pq.alloc()
     pq.inst(Y(2), X(qq))
     assert pq.get_qubits() == {0, 1, 2, 4, 5}  # this synthesizes the allocation
+
+    qubit_index = 1
+    p = Program(("H", qubit_index))
+    assert p.get_qubits() == {qubit_index}
+    q1 = p.alloc()
+    q2 = p.alloc()
+    p.inst(("CNOT", q1, q2))
+    assert p.get_qubits() == {qubit_index, q1.index(), q2.index()}