Improved Features types

bofire/benchmarks/benchmark.py

@@ -135,31 +135,28 @@ def __init__(
     ):
         super().__init__(**kwargs)
         self._benchmark = benchmark
+        benchmark_inputs = self._benchmark.domain.inputs.get()
         assert n_filler_features >= 1, "n_filler_features must be >= 1."
         assert len(benchmark.domain.constraints) == 0, "Constraints not supported yet."
-        assert len(benchmark.domain.inputs.get(ContinuousInput)) == len(
-            benchmark.domain.inputs
-        ), "Only continuous inputs supported yet."
+        if not Inputs.is_continuous(benchmark_inputs):
+            raise ValueError("Only continuous inputs supported yet.")
         self.n_filler_features = n_filler_features
         self.n_features_per_original_feature = n_features_per_original_feature
 
         features = []
         constraints = []
-        for j, feat in enumerate(self._benchmark.domain.inputs.get()):
+        for j, feat in enumerate(benchmark_inputs):
             features += [
                 ContinuousInput(
                     key=f"{feat.key}_{i}",
-                    bounds=(0, 1 / len(self._benchmark.domain.inputs)),
+                    bounds=(0, 1 / len(benchmark_inputs)),
                 )
                 if self.n_features_per_original_feature == 1
                 else ContinuousDescriptorInput(
                     key=f"{feat.key}_{i}",
                     bounds=(0, 1 / len(self._benchmark.domain.inputs)),
                     descriptors=self._benchmark.domain.inputs.get_keys(),
-                    values=[
-                        1 if k == j else 0
-                        for k in range(len(self._benchmark.domain.inputs))
-                    ],
+                    values=[1 if k == j else 0 for k in range(len(benchmark_inputs))],
                 )
                 for i in range(self.n_features_per_original_feature)
             ]
@@ -171,7 +168,7 @@ def __init__(
                             for i in range(self.n_features_per_original_feature)
                         ],
                         coefficients=[1.0] * self.n_features_per_original_feature,
-                        rhs=1 / len(self._benchmark.domain.inputs),
+                        rhs=1 / len(benchmark_inputs),
                     )
                 )
 
@@ -207,18 +204,13 @@ def __init__(
             constraints=Constraints(constraints=constraints),
             outputs=self._benchmark.domain.outputs,
         )
-        self._mins = np.array(
-            [feat.bounds[0] for feat in self._benchmark.domain.inputs.get()]
-        )
+
+        self._mins = np.array([feat.bounds[0] for feat in benchmark_inputs])
         self._scales = np.array(
-            [
-                feat.bounds[1] - feat.bounds[0]
-                for feat in self._benchmark.domain.inputs.get()
-            ]
+            [feat.bounds[1] - feat.bounds[0] for feat in benchmark_inputs]
         )
         self._scales_new = np.array(
-            [1 / len(self._benchmark.domain.inputs.get_keys())]
-            * len(self._benchmark.domain.inputs.get_keys())
+            [1 / len(benchmark_inputs.get_keys())] * len(benchmark_inputs.get_keys())
         )
 
     def _transform(self, X: pd.DataFrame) -> pd.DataFrame:

bofire/data_models/domain/constraints.py

@@ -1,38 +1,44 @@
 import collections.abc
 from collections.abc import Iterator, Sequence
 from itertools import chain
-from typing import Generic, List, Literal, Optional, Type, TypeVar, Union
+from typing import Generic, Literal, Type, Union, overload
 
 import pandas as pd
 from pydantic import Field
+from typing_extensions import Self, TypeVar
 
 from bofire.data_models.base import BaseModel
 from bofire.data_models.constraints.api import AnyConstraint, Constraint
 from bofire.data_models.filters import filter_by_class
 
 
-C = TypeVar("C", bound=Union[AnyConstraint, Constraint])
-CIncludes = TypeVar("CIncludes", bound=Union[AnyConstraint, Constraint])
-CExcludes = TypeVar("CExcludes", bound=Union[AnyConstraint, Constraint])
+ConstraintT = TypeVar(
+    "ConstraintT", bound=AnyConstraint | Constraint, default=AnyConstraint | Constraint
+)
+ConstraintGetT = TypeVar(
+    "ConstraintGetT",
+    bound=AnyConstraint | Constraint,
+    default=AnyConstraint | Constraint,
+)
 
 
-class Constraints(BaseModel, Generic[C]):
+class Constraints(BaseModel, Generic[ConstraintT]):
     type: Literal["Constraints"] = "Constraints"
-    constraints: Sequence[C] = Field(default_factory=list)
+    constraints: Sequence[ConstraintT] = Field(default_factory=list)
 
-    def __iter__(self) -> Iterator[C]:
+    def __iter__(self) -> Iterator[ConstraintT]:
         return iter(self.constraints)
 
     def __len__(self):
         return len(self.constraints)
 
-    def __getitem__(self, i) -> C:
+    def __getitem__(self, i) -> ConstraintT:
         return self.constraints[i]
 
     def __add__(
         self,
-        other: Union[Sequence[CIncludes], "Constraints[CIncludes]"],
-    ) -> "Constraints[Union[C, CIncludes]]":
+        other: Union[Sequence[ConstraintGetT], "Constraints[ConstraintGetT]"],
+    ) -> "Constraints[Union[ConstraintT, ConstraintGetT]]":
         if isinstance(other, collections.abc.Sequence):
             other_constraints = other
         else:
@@ -87,14 +93,28 @@ def is_fulfilled(self, experiments: pd.DataFrame, tol: float = 1e-6) -> pd.Serie
             .all(axis=1)
         )
 
+    @overload
     def get(
         self,
-        includes: Union[
-            Type[CIncludes], Sequence[Type[CIncludes]]
-        ] = Constraint,  # ty: ignore[invalid-parameter-default]
-        excludes: Optional[Union[Type[CExcludes], List[Type[CExcludes]]]] = None,
+        includes: Type[ConstraintGetT] | Sequence[Type[ConstraintGetT]],
+        excludes: None = None,
         exact: bool = False,
-    ) -> "Constraints[CIncludes]":
+    ) -> "Constraints[ConstraintGetT]": ...
+
+    @overload
+    def get(
+        self,
+        includes: Type[ConstraintGetT] | Sequence[Type[ConstraintGetT]] = Constraint,
+        excludes: Type[AnyConstraint] | Sequence[Type[AnyConstraint]] | None = None,
+        exact: bool = False,
+    ) -> Self: ...
+
+    def get(
+        self,
+        includes: Type[ConstraintGetT] | Sequence[Type[ConstraintGetT]] = Constraint,
+        excludes: Type[AnyConstraint] | Sequence[Type[AnyConstraint]] | None = None,
+        exact: bool = False,
+    ) -> "Constraints[ConstraintGetT]":
         """Get constraints of the domain
 
         Args: