dtflow 0.4.0__tar.gz → 0.4.1__tar.gz

{dtflow-0.4.0 → dtflow-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dtflow
-Version: 0.4.0
+Version: 0.4.1
 Summary: A flexible data transformation tool for ML training formats (SFT, RLHF, Pretrain)
 Project-URL: Homepage, https://github.com/yourusername/DataTransformer
 Project-URL: Documentation, https://github.com/yourusername/DataTransformer#readme

{dtflow-0.4.0 → dtflow-0.4.1}/dtflow/__init__.py

@@ -42,7 +42,7 @@ from .tokenizers import (
     token_stats,
 )
 
-__version__ = "0.4.0"
+__version__ = "0.4.1"
 
 __all__ = [
     # core

{dtflow-0.4.0 → dtflow-0.4.1}/dtflow/cli/commands.py

@@ -796,6 +796,17 @@ def _generate_default_transform(field_names: List[str]) -> str:
     return "\n".join(lines) if lines else "        # 在这里定义输出字段"
 
 
+def _unwrap(obj: Any) -> Any:
+    """递归将 DictWrapper 转换为普通 dict"""
+    if hasattr(obj, "to_dict"):
+        return _unwrap(obj.to_dict())
+    if isinstance(obj, dict):
+        return {k: _unwrap(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [_unwrap(v) for v in obj]
+    return obj
+
+
 def _execute_transform(
     input_path: Path,
     config_path: Path,
@@ -829,7 +840,8 @@ def _execute_transform(
         try:
             # 包装转换函数以支持属性访问（配置文件中定义的 Item 类）
             def wrapped_transform(item):
-                return transform_func(DictWrapper(item))
+                result = transform_func(DictWrapper(item))
+                return _unwrap(result)
 
             st = load_stream(str(input_path))
             if num:
@@ -926,7 +938,8 @@ def _execute_preset_transform(
         try:
             # 包装转换函数以支持属性访问
             def wrapped_transform(item):
-                return transform_func(DictWrapper(item))
+                result = transform_func(DictWrapper(item))
+                return _unwrap(result)
 
             st = load_stream(str(input_path))
             if num:

{dtflow-0.4.0 → dtflow-0.4.1}/dtflow/streaming.py

@@ -84,6 +84,8 @@ class StreamingTransformer:
         self._source_path = source_path
         self._total = total
         self._operations: List[Dict[str, Any]] = []
+        self._error_count = 0
+        self._first_error: Optional[str] = None
 
     @classmethod
     def load_stream(cls, filepath: str, batch_size: int = 10000) -> "StreamingTransformer":
@@ -194,17 +196,20 @@ class StreamingTransformer:
         Returns:
             新的 StreamingTransformer（惰性，不立即执行）
         """
+        # transform 是 1:1 转换，保留 total
+        new_st = StreamingTransformer(iter([]), self._source_path, total=self._total)
+        new_st._operations = self._operations + [{"type": "transform", "func": func}]
 
         def transformed_iterator():
             for item in self._iterator:
                 try:
                     yield func(item)
-                except Exception:
-                    pass  # 跳过错误
+                except Exception as e:
+                    new_st._error_count += 1
+                    if new_st._first_error is None:
+                        new_st._first_error = f"{type(e).__name__}: {e}"
 
-        # transform 是 1:1 转换，保留 total
-        new_st = StreamingTransformer(transformed_iterator(), self._source_path, total=self._total)
-        new_st._operations = self._operations + [{"type": "transform", "func": func}]
+        new_st._iterator = transformed_iterator()
         return new_st
 
     def head(self, n: int) -> "StreamingTransformer":
@@ -299,16 +304,21 @@ class StreamingTransformer:
         ext = path.suffix.lower()
 
         if ext == ".jsonl":
-            return self._save_jsonl(filepath, show_progress)
+            count = self._save_jsonl(filepath, show_progress)
         elif ext == ".csv":
-            return self._save_batched(filepath, "csv", batch_size, show_progress)
+            count = self._save_batched(filepath, "csv", batch_size, show_progress)
         elif ext == ".parquet":
-            return self._save_batched(filepath, "parquet", batch_size, show_progress)
+            count = self._save_batched(filepath, "parquet", batch_size, show_progress)
         elif ext in (".arrow", ".feather"):
-            return self._save_batched(filepath, "arrow", batch_size, show_progress)
+            count = self._save_batched(filepath, "arrow", batch_size, show_progress)
         else:
-            # 默认 JSONL
-            return self._save_jsonl(filepath, show_progress)
+            count = self._save_jsonl(filepath, show_progress)
+
+        # 打印错误摘要
+        if self._error_count > 0:
+            print(f"⚠️  跳过 {self._error_count} 条错误记录: {self._first_error}")
+
+        return count
 
     def _save_jsonl(self, filepath: str, show_progress: bool) -> int:
         """JSONL 逐行流式保存（使用 orjson）"""

{dtflow-0.4.0 → dtflow-0.4.1}/tests/test_streaming.py

@@ -279,3 +279,22 @@ class TestEdgeCases:
         items2 = st.collect()
         assert len(items2) == 0
         os.unlink(temp_jsonl)
+
+    def test_transform_error_tracking(self, temp_jsonl, tmp_path):
+        """测试转换错误跟踪"""
+        st = load_stream(temp_jsonl)
+
+        def bad_transform(x):
+            if x["id"] % 10 == 5:  # 5, 15, 25, ... 共 10 条会出错
+                raise KeyError("missing_key")
+            return {"new_id": x["id"]}
+
+        transformed = st.transform(bad_transform)
+        output_path = tmp_path / "output.jsonl"
+        count = transformed.save(str(output_path), show_progress=False)
+
+        # 验证结果
+        assert count == 90  # 100 - 10 = 90
+        assert transformed._error_count == 10
+        assert "KeyError" in transformed._first_error
+        os.unlink(temp_jsonl)

{dtflow-0.4.0 → dtflow-0.4.1}/tests/test_transformer.py

@@ -610,5 +610,54 @@ class TestConcat:
         assert len(dt2) == 1
 
 
+class TestUnwrap:
+    """测试 _unwrap 函数（DictWrapper 转 dict）"""
+
+    def test_unwrap_simple(self):
+        """测试简单 DictWrapper 转换"""
+        from dtflow.cli.commands import _unwrap
+
+        wrapper = DictWrapper({"a": 1, "b": "text"})
+        result = _unwrap(wrapper)
+
+        assert result == {"a": 1, "b": "text"}
+        assert type(result) is dict
+
+    def test_unwrap_nested(self):
+        """测试嵌套 DictWrapper 转换"""
+        from dtflow.cli.commands import _unwrap
+
+        data = {"outer": {"inner": {"value": 123}}}
+        wrapper = DictWrapper(data)
+        # 访问嵌套会产生新的 DictWrapper
+        nested_wrapper = wrapper.outer.inner
+
+        result = _unwrap({"data": nested_wrapper, "list": [wrapper.outer]})
+
+        assert result == {"data": {"value": 123}, "list": [{"inner": {"value": 123}}]}
+        assert type(result["data"]) is dict
+        assert type(result["list"][0]) is dict
+
+    def test_unwrap_in_list(self):
+        """测试列表中的 DictWrapper 转换"""
+        from dtflow.cli.commands import _unwrap
+
+        wrapper = DictWrapper({"x": 1})
+        result = _unwrap([wrapper, {"y": 2}, wrapper])
+
+        assert result == [{"x": 1}, {"y": 2}, {"x": 1}]
+        assert all(type(item) is dict for item in result)
+
+    def test_unwrap_plain_dict(self):
+        """测试普通 dict 不受影响"""
+        from dtflow.cli.commands import _unwrap
+
+        data = {"a": 1, "nested": {"b": 2}}
+        result = _unwrap(data)
+
+        assert result == data
+        assert type(result) is dict
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])