# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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.
"""AutoQASM types and type utilities."""
from __future__ import annotations
from collections.abc import Iterable
from typing import Any, get_args
import numpy as np
import oqpy
import oqpy.base
from openpulse import ast
from autoqasm import errors, program
from braket.circuits import FreeParameterExpression
from braket.registers import Qubit
[docs]
def is_qasm_type(val: Any) -> bool:
"""Returns whether the provided object is of a QASM type or is itself a QASM type
which receives special treatment by AutoQASM.
Args:
val (Any): The object of which to check the type.
Returns:
bool: Whether the object is a QASM type.
"""
qasm_types = (
oqpy.Range,
oqpy._ClassicalVar,
oqpy.Qubit,
oqpy.base.OQPyExpression,
FreeParameterExpression,
)
# The input can either be a class, like oqpy.Range ...
if type(val) is type:
return issubclass(val, qasm_types)
# ... or an instance of a class, like oqpy.Range(10)
return isinstance(val, qasm_types)
[docs]
def make_annotations_list(annotations: str | Iterable[str] | None) -> list[str]:
return [annotations] if isinstance(annotations, str) else annotations or []
QubitIdentifierType = int | str | Qubit | oqpy._ClassicalVar | oqpy.base.OQPyExpression | oqpy.Qubit
# Precompute the type tuple once: ``get_args(QubitIdentifierType)`` isn't
# free, and this is called on every gate emission.
_QUBIT_IDENTIFIER_TYPES: tuple[type, ...] = get_args(QubitIdentifierType)
[docs]
def is_qubit_identifier_type(qubit: Any) -> bool:
"""Checks if a given object is a qubit identifier type.
Args:
qubit (Any): The object to check.
Returns:
bool: True if the object is a qubit identifier type, False otherwise.
"""
return isinstance(qubit, _QUBIT_IDENTIFIER_TYPES)
[docs]
class Range(oqpy.Range):
def __init__(
self,
start: int,
stop: int | None = None,
step: int | None = 1,
annotations: str | Iterable[str] | None = None,
):
"""Creates a range definition.
Args:
start (int): Start of the range.
stop (int | None): End of the range. Defaults to None.
step (int | None): Step of the range. Defaults to 1.
annotations (str | Iterable[str] | None): Annotations for the range.
"""
if stop is None:
stop = start
start = 0
super().__init__(start, stop, step)
self.annotations = make_annotations_list(annotations)
[docs]
class ArrayVar(oqpy.ArrayVar):
def __init__(
self,
init_expression: Iterable,
*args,
annotations: str | Iterable[str] | None = None,
**kwargs,
):
if (
program.get_program_conversion_context().subroutines_processing
or not program.get_program_conversion_context().at_function_root_scope
):
raise errors.InvalidArrayDeclaration(
"Arrays may only be declared at the root scope of an AutoQASM main function."
)
if not isinstance(init_expression, Iterable):
raise errors.InvalidArrayDeclaration("init_expression must be an iterable type.")
dimensions = np.shape(init_expression)
super().__init__(
init_expression=init_expression,
*args, # noqa: B026
annotations=make_annotations_list(annotations),
dimensions=dimensions,
**kwargs,
)
self.name = program.get_program_conversion_context().next_var_name(oqpy.ArrayVar)
[docs]
class BitVar(oqpy.BitVar):
def __init__(self, *args, annotations: str | Iterable[str] | None = None, **kwargs):
super().__init__(*args, annotations=make_annotations_list(annotations), **kwargs)
self.name = program.get_program_conversion_context().next_var_name(oqpy.BitVar)
if self.size and self.init_expression != "output":
value = self.init_expression or 0
self.init_expression = ast.BitstringLiteral(value, self.size)
[docs]
class BoolVar(oqpy.BoolVar):
def __init__(self, *args, annotations: str | Iterable[str] | None = None, **kwargs):
super().__init__(*args, annotations=make_annotations_list(annotations), **kwargs)
self.name = program.get_program_conversion_context().next_var_name(oqpy.BoolVar)
[docs]
class FloatVar(oqpy.FloatVar):
def __init__(self, *args, annotations: str | Iterable[str] | None = None, **kwargs):
super().__init__(*args, annotations=make_annotations_list(annotations), **kwargs)
self.name = program.get_program_conversion_context().next_var_name(oqpy.FloatVar)
[docs]
class IntVar(oqpy.IntVar):
def __init__(self, *args, annotations: str | Iterable[str] | None = None, **kwargs):
super().__init__(*args, annotations=make_annotations_list(annotations), **kwargs)
self.name = program.get_program_conversion_context().next_var_name(oqpy.IntVar)