[feat] Add validation for func output

Author: yhteoh

src/oqd_dataschema/datastore.py

@@ -18,7 +18,7 @@
 
 import json
 import pathlib
-from typing import Any, Dict, Literal
+from typing import Any, Callable, Dict, Literal
 
 import h5py
 from pydantic import (
@@ -220,8 +220,16 @@ def add(self, **groups):
 
         self.update(**groups)
 
-    def pipe(self, func) -> Datastore:
-        return func(self)
+    def pipe(self, func: Callable[[Datastore], Datastore]) -> Datastore:
+        _result = func(self)
+
+        if _result is None:
+            return self
+
+        if _result is not Datastore:
+            raise ValueError("`func` must return None or Datastore.")
+
+        return _result
 
 
 # %%