From 433040fbe146f02273853e90321fc2e6aad3d9a0 Mon Sep 17 00:00:00 2001 From: Serwan Asaad Date: Thu, 12 Dec 2024 11:58:20 +0100 Subject: [PATCH] working measure operation --- quam/core/operation/function_properties.py | 47 ++++++------------- .../quantum_components/test_qubit.py | 8 +++- tests/operations/test_function_properties.py | 42 +++++++++++++++-- tests/operations/test_operations.py | 12 +++++ 4 files changed, 69 insertions(+), 40 deletions(-) diff --git a/quam/core/operation/function_properties.py b/quam/core/operation/function_properties.py index 7449dc8..1a1c343 100644 --- a/quam/core/operation/function_properties.py +++ b/quam/core/operation/function_properties.py @@ -17,10 +17,10 @@ class FunctionProperties: """ Properties of a quantum operation function. - + This class extracts and stores metadata about functions that operate on quantum components, including argument information and type requirements. - + Attributes: quantum_component_name: Name of the parameter accepting the quantum component quantum_component_type: Type of quantum component the function operates on @@ -28,17 +28,19 @@ class FunctionProperties: required_args: List of required argument names after the quantum component optional_args: Dictionary of optional arguments and their default values """ + quantum_component_name: str quantum_component_type: Type[QC] name: str = "" required_args: list[str] = field(default_factory=list) optional_args: dict[str, Any] = field(default_factory=dict) + return_type: Optional[Type] = None def __post_init__(self): # Make a new list/dict to avoid sharing between instances self.required_args = list(self.required_args) self.optional_args = dict(self.optional_args) - + # Validate argument names all_args = self.required_args + list(self.optional_args) for arg in all_args: @@ -47,31 +49,6 @@ def __post_init__(self): if keyword.iskeyword(arg): raise ValueError(f"Argument name cannot be a Python keyword: {arg!r}") - @staticmethod - def _resolve_type(type_hint: Any) -> Optional[Type]: - """ - Resolve type hints, including string forward references and complex types. - - Args: - type_hint: Any type annotation - - Returns: - The resolved base type, or None if unresolvable - """ - if type_hint is None: - return None - # Handle string forward references - if isinstance(type_hint, str): - return None - # Handle Optional, Union, etc - if get_origin(type_hint) is not None: - args = get_args(type_hint) - return args[0] if args else None - # Handle regular types - if isinstance(type_hint, type): - return type_hint - return None - @staticmethod def _is_quantum_component_type(type_hint: Optional[Type]) -> bool: """Check if type is or inherits from QuantumComponent.""" @@ -101,15 +78,16 @@ def from_function(cls, func: Callable) -> "FunctionProperties": # Try to get type hints, gracefully handle missing annotations try: type_hints = get_type_hints(func) - except Exception: - type_hints = {} + except (NameError, TypeError): + # Fallback to using the raw annotations if get_type_hints fails + type_hints = getattr(func, "__annotations__", {}) parameters_iterator = iter(parameters) first_param_name = next(parameters_iterator) - + # Get and resolve the type of the first parameter - first_param_type = cls._resolve_type(type_hints.get(first_param_name)) - + first_param_type = type_hints.get(first_param_name) + if not cls._is_quantum_component_type(first_param_type): if first_param_type is None: msg = ( @@ -138,4 +116,7 @@ def from_function(cls, func: Callable) -> "FunctionProperties": # Store the default value directly function_properties.optional_args[param_name] = param.default + # Get the return type from the function annotations + function_properties.return_type = type_hints.get("return") + return function_properties diff --git a/tests/components/quantum_components/test_qubit.py b/tests/components/quantum_components/test_qubit.py index e6704ef..f5c8ba3 100644 --- a/tests/components/quantum_components/test_qubit.py +++ b/tests/components/quantum_components/test_qubit.py @@ -63,7 +63,9 @@ def test_qubit_align(mock_qubit_with_resonator, mock_qubit, mocker): from quam.components.quantum_components.qubit import align - align.assert_called_once_with(*{"q1.xy", "q1.resonator", "q0.xy"}) + align.assert_called_once() + called_args, _ = align.call_args + assert set(called_args) == {"q1.xy", "q1.resonator", "q0.xy"} def test_qubit_get_macros(mock_qubit): @@ -77,7 +79,9 @@ def test_qubit_apply_align(mock_qubit_with_resonator, mocker): from quam.components.quantum_components.qubit import align - align.assert_called_once_with(*{"q1.xy", "q1.resonator"}) + align.assert_called_once() + called_args, _ = align.call_args + assert set(called_args) == {"q1.xy", "q1.resonator"} def test_qubit_inferred_id_direct(): diff --git a/tests/operations/test_function_properties.py b/tests/operations/test_function_properties.py index 9142680..0e4c307 100644 --- a/tests/operations/test_function_properties.py +++ b/tests/operations/test_function_properties.py @@ -220,11 +220,43 @@ def operation(component, arg1, arg2=None): FunctionProperties.from_function(operation) -def test_from_function_with_non_type_annotation(): - """Test handling of invalid type annotations.""" +def test_from_function_with_return_type(): + """Test that return type is correctly captured.""" - def operation(component: "not a real type", arg1: int): # type: ignore + def operation(component: DummyQuantumComponent) -> int: + return 42 + + props = FunctionProperties.from_function(operation) + assert props.return_type == int + + +def test_from_function_with_optional_return_type(): + """Test handling of Optional return type.""" + from typing import Optional + + def operation(component: DummyQuantumComponent) -> Optional[int]: + return None + + props = FunctionProperties.from_function(operation) + assert props.return_type == Optional[int] + + +def test_from_function_with_qua_return_type(): + """Test handling of QUA variable return types.""" + from qm.qua._expressions import QuaBoolType + + def operation(component: DummyQuantumComponent) -> QuaBoolType: pass - with pytest.raises(ValueError, match="missing type annotation"): - FunctionProperties.from_function(operation) + props = FunctionProperties.from_function(operation) + assert props.return_type == QuaBoolType + + +def test_from_function_without_return_type(): + """Test handling of functions without return type annotation.""" + + def operation(component: DummyQuantumComponent): + pass + + props = FunctionProperties.from_function(operation) + assert props.return_type is None diff --git a/tests/operations/test_operations.py b/tests/operations/test_operations.py index af86c24..4763d3d 100644 --- a/tests/operations/test_operations.py +++ b/tests/operations/test_operations.py @@ -108,3 +108,15 @@ def test_op(qubit: Qubit, arg1: float, arg2: str = "default"): assert result[0] == test_qubit assert result[1] == (1.0,) # arg1 as positional arg assert result[2] == {"arg2": "test"} # arg2 as kwarg + + +def test_measure_operation(test_qubit): + from qm.qua._expressions import QuaBoolType + + def measure(qubit: Qubit, **kwargs) -> QuaBoolType: + pass + + op = Operation(measure) + + assert op.properties.return_type == QuaBoolType + assert op.func == measure