refactoring to use annotations

src/datamaestro/data/ml.py

@@ -1,19 +1,27 @@
 """Machine learning generic data formats"""
-from typing import List
+from typing import Generic, TypeVar, Optional
 from pathlib import Path
-from . import Base, argument
+from experimaestro import Param, Meta, argument
+from . import Base
 
+Train = TypeVar("Train", bound=Base)
+Validation = TypeVar("Validation", bound=Base)
+Test = TypeVar("Test", bound=Base)
 
-@argument("train", type=Base, help="The training dataset")
-@argument("validation", type=Base, help="The validation dataset", required=False)
-@argument("test", type=Base, help="The test dataset", required=False)
-class Supervised(Base):
-    pass
+
+class Supervised(Base, Generic[Train, Validation, Test]):
+    train: Param[Base]
+    """The training dataset"""
+
+    validation: Param[Optional[Base]] = None
+    """The validation dataset (optional)"""
+
+    test: Param[Optional[Base]] = None
+    """The training optional"""
 
 
-@argument("path", type=Path)
 @argument("classes")
 class FolderBased(Base):
     """Classification dataset where folders give the basis"""
 
-    pass
+    path: Meta[Path]