From 45df6c6caa2b85fac0acc2f573a7aebbe41d8e56 Mon Sep 17 00:00:00 2001
From: Ryan Shaffer <3620100+rmshaffer@users.noreply.github.com>
Date: Tue, 3 Oct 2023 16:56:55 -0400
Subject: [PATCH] Upgrade to oqpy 0.3.2

---
 setup.py                                              |  5 +----
 src/braket/experimental/autoqasm/types/conversions.py | 11 +++++++++--
 .../braket/experimental/autoqasm/test_types.py        |  4 ++--
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/setup.py b/setup.py
index a206df5d1..04be9bd8d 100644
--- a/setup.py
+++ b/setup.py
@@ -33,10 +33,7 @@
         # simulation of mid-circuit measurement, which AutoQASM requires.
         # NOTE: This change should remain in the feature/autoqasm branch; do not merge to main.
         "amazon-braket-default-simulator @ git+https://github.com/aws/amazon-braket-default-simulator-python.git@31d6c95f3ac250a0ccd04e1433ad61c8bfa4bde4#egg=amazon-braket-default-simulator",  # noqa E501
-        # Pin the latest commit of the main branch of openqasm/oqpy.git to get the version of
-        # oqpy which contains changes that AutoQASM relies on.
-        # NOTE: This change should remain in the feature/autoqasm branch; do not merge to main.
-        "oqpy @ git+https://github.com/openqasm/oqpy.git@4bb47dd59e07f03234872b1a5ec5c3c767abbdc1#egg=oqpy",  # noqa E501
+        "oqpy~=0.3.2",
         "setuptools",
         "backoff",
         "boltons",
diff --git a/src/braket/experimental/autoqasm/types/conversions.py b/src/braket/experimental/autoqasm/types/conversions.py
index 61b7c0c68..e78b3dff0 100644
--- a/src/braket/experimental/autoqasm/types/conversions.py
+++ b/src/braket/experimental/autoqasm/types/conversions.py
@@ -14,7 +14,7 @@
 """Type conversions between Python and the autoqasm representation for types."""
 
 import typing
-from functools import singledispatch
+from functools import partialmethod, singledispatch
 from typing import Any, Union
 
 import numpy as np
@@ -52,6 +52,13 @@ def map_type(python_type: type) -> type:
             raise errors.ParameterTypeError(
                 f"Unsupported array type: {item_type}. AutoQASM arrays only support ints."
             )
+
+        def _partial_class(cls, *args, **kwargs) -> type:
+            class PartialCls(cls):
+                __init__ = partialmethod(cls.__init__, *args, **kwargs)
+
+            return PartialCls
+
         # TODO: Update array length to match the input rather than hardcoding
         # OQPY and QASM require arrays have a set length. python doesn't require this,
         # so the length of the array is indeterminate.
@@ -60,7 +67,7 @@ def map_type(python_type: type) -> type:
         # Here's where the info is stored for oqpy variables:
         # ctx = program.get_program_conversion_context()
         # dims = ctx.get_oqpy_program().declared_vars[name_of_var].dimensions
-        return oqpy.ArrayVar[oqpy.IntVar, 10]
+        return _partial_class(oqpy.ArrayVar, dimensions=[10], base_type=oqpy.IntVar)
     if issubclass(origin_type, tuple):
         raise TypeError(
             "Tuples are not supported as parameters to AutoQASM functions; "
diff --git a/test/unit_tests/braket/experimental/autoqasm/test_types.py b/test/unit_tests/braket/experimental/autoqasm/test_types.py
index 834fc54ec..c5b975773 100644
--- a/test/unit_tests/braket/experimental/autoqasm/test_types.py
+++ b/test/unit_tests/braket/experimental/autoqasm/test_types.py
@@ -207,7 +207,7 @@ def test_return_python_array():
     """Test returning a python array of ints."""
 
     @aq.subroutine
-    def tester(arr: List[int]) -> List[int]:
+    def tester() -> List[int]:
         return [1, 2, 3]
 
     @aq.main(num_qubits=4)
@@ -215,7 +215,7 @@ def main():
         tester()
 
     expected = """OPENQASM 3.0;
-def tester(array[int[32], 10] arr) -> array[int[32], 10] {
+def tester() -> array[int[32], 10] {
     array[int[32], 10] retval_;
     retval_ = {1, 2, 3};
     return retval_;