Fix lint, update tests with doc building

horqrux/abstract.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Callable, Iterable, Tuple
+from typing import Any, Iterable, Tuple
 
 import numpy as np
 from jax import Array
@@ -48,7 +48,7 @@ def __post_init__(self) -> None:
     def __iter__(self) -> Iterable:
         return iter((self.generator_name, self.target, self.control))
 
-    def tree_flatten(self) -> Tuple[Tuple, Tuple[str, QubitSupport, QubitSupport]]:
+    def tree_flatten(self) -> Tuple[Tuple, Tuple[str, TargetQubits, ControlQubits]]:
         children = ()
         aux_data = (self.generator_name, self.target, self.control)
         return (children, aux_data)
@@ -86,13 +86,14 @@ def __post_init__(self) -> None:
         super().__post_init__()
 
         def parse_dict(values: dict[str, float] = {}) -> float:
-            return values[self.param]
+            return values[self.param]  # type: ignore[index]
 
-        self.parse_values: Callable[[dict[str, float]], float] = (
-            parse_dict if isinstance(self.param, str) else lambda x: self.param
-        )
+        def parse_val(values: dict[str, float] = {}) -> float:
+            return self.param  # type: ignore[return-value]
+
+        self.parse_values = parse_dict if isinstance(self.param, str) else parse_val
 
-    def tree_flatten(self) -> Tuple[Tuple, Tuple[str, TargetQubits, ControlQubits, str | float]]:
+    def tree_flatten(self) -> Tuple[Tuple, Tuple[str, Tuple, Tuple, str | float]]:  # type: ignore[override]
         children = ()
         aux_data = (
             self.name,

horqrux/parametric.py

@@ -18,7 +18,7 @@ def RX(param: float | str, target: TargetQubits, control: ControlQubits = (None,
     Returns:
         Parametric: A Parametric gate object.
     """
-    return Parametric("X", target, control, param=param)
+    return Parametric("X", target, control, param)
 
 
 def RY(param: float | str, target: TargetQubits, control: ControlQubits = (None,)) -> Parametric:
@@ -32,7 +32,7 @@ def RY(param: float | str, target: TargetQubits, control: ControlQubits = (None,
     Returns:
         Parametric: A Parametric gate object.
     """
-    return Parametric("Y", target, control, param=param)
+    return Parametric("Y", target, control, param)
 
 
 def RZ(param: float | str, target: TargetQubits, control: ControlQubits = (None,)) -> Parametric:
@@ -46,7 +46,7 @@ def RZ(param: float | str, target: TargetQubits, control: ControlQubits = (None,
     Returns:
         Parametric: A Parametric gate object.
     """
-    return Parametric("Z", target, control, param=param)
+    return Parametric("Z", target, control, param)
 
 
 class _PHASE(Parametric):

pyproject.toml

@@ -43,7 +43,7 @@ features = [
 ]
 
 [tool.hatch.envs.tests.scripts]
-test = "pytest -n auto {args}"
+test = "pytest -n auto {args} && hatch -e docs run docs:build"
 
 [tool.hatch.envs.docs]
 dependencies = [