Move feature classes to features.py

miplearn/features.py

@@ -4,20 +4,57 @@
 
 import collections
 import numbers
-from typing import TYPE_CHECKING, Dict
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, Optional, Set, List, Hashable
 
-from miplearn.types import (
-    Features,
-    ConstraintFeatures,
-    InstanceFeatures,
-    VariableFeatures,
-    VarIndex,
-)
+from miplearn.types import VarIndex, Solution
 
 if TYPE_CHECKING:
     from miplearn import InternalSolver, Instance
 
 
+@dataclass
+class TrainingSample:
+    lp_log: Optional[str] = None
+    lp_solution: Optional[Solution] = None
+    lp_value: Optional[float] = None
+    lazy_enforced: Optional[Set[str]] = None
+    lower_bound: Optional[float] = None
+    mip_log: Optional[str] = None
+    solution: Optional[Solution] = None
+    upper_bound: Optional[float] = None
+    slacks: Optional[Dict[str, float]] = None
+
+
+@dataclass
+class InstanceFeatures:
+    user_features: Optional[List[float]] = None
+    lazy_constraint_count: int = 0
+
+
+@dataclass
+class VariableFeatures:
+    category: Optional[Hashable] = None
+    user_features: Optional[List[float]] = None
+
+
+@dataclass
+class ConstraintFeatures:
+    rhs: Optional[float] = None
+    lhs: Optional[Dict[str, float]] = None
+    sense: Optional[str] = None
+    category: Optional[Hashable] = None
+    user_features: Optional[List[float]] = None
+    lazy: bool = False
+
+
+@dataclass
+class Features:
+    instance: Optional[InstanceFeatures] = None
+    variables: Optional[Dict[str, Dict[VarIndex, VariableFeatures]]] = None
+    constraints: Optional[Dict[str, ConstraintFeatures]] = None
+
+
 class FeaturesExtractor:
     def __init__(
         self,