working measure operation

quam/core/operation/function_properties.py

@@ -17,28 +17,30 @@
 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
         name: Name of the function
         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

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():

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

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