fix tests

epoch8 · Aug 19, 2024 · 181cf40 · 181cf40
1 parent 0008bf9
commit 181cf40
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/datapipe/step/batch_transform.py b/datapipe/step/batch_transform.py
@@ -546,8 +546,6 @@ def _apply_filters_to_run_config(
                 dt = ds.get_table(self.filters)
                 df = dt.get_data()
                 filters = cast(List[LabelDict], df[dt.primary_keys].to_dict(orient="records"))
-            elif isinstance(self.filters, DataTable):
-                filters = cast(List[LabelDict], self.filters.get_data().to_dict(orient="records"))
             elif isinstance(self.filters, pd.DataFrame):
                 filters = cast(List[LabelDict], self.filters.to_dict(orient="records"))
             elif isinstance(self.filters, list) and all([isinstance(x, dict) for x in self.filters]):

diff --git a/datapipe/types.py b/datapipe/types.py
@@ -42,7 +42,7 @@
 TransformResult = Union[DataDF, List[DataDF], Tuple[DataDF, ...]]
 
 LabelDict = Dict[str, Any]
-Filters = Union[str, "DataTable", IndexDF, List[LabelDict], Callable[..., List[LabelDict]], Callable[..., IndexDF]]
+Filters = Union[str, IndexDF, List[LabelDict], Callable[..., List[LabelDict]], Callable[..., IndexDF]]
 try:
     from sqlalchemy.orm import DeclarativeBase