dtflow 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

dtflow/core.py

@@ -6,29 +6,16 @@ DataTransformer 核心模块
 from typing import List, Dict, Any, Optional, Callable, Union, Tuple, Literal
 from copy import deepcopy
 from dataclasses import dataclass
-import json
 
-from .storage.io import save_data, load_data
+import orjson
 
-# 尝试使用 orjson（更快的 JSON 序列化库）
-try:
-    import orjson
-    _HAS_ORJSON = True
-except ImportError:
-    _HAS_ORJSON = False
+from .storage.io import save_data, load_data
+from .lineage import LineageTracker
 
 
 def _fast_json_dumps(obj: Any) -> str:
-    """
-    快速 JSON 序列化，优先使用 orjson。
-
-    orjson 比标准 json 快约 10 倍，特别适合大量数据的序列化场景。
-    """
-    if _HAS_ORJSON:
-        # orjson.dumps 返回 bytes，需要 decode
-        return orjson.dumps(obj, option=orjson.OPT_SORT_KEYS).decode('utf-8')
-    else:
-        return json.dumps(obj, sort_keys=True, ensure_ascii=False)
+    """快速 JSON 序列化（使用 orjson，比标准 json 快约 10 倍）"""
+    return orjson.dumps(obj, option=orjson.OPT_SORT_KEYS).decode('utf-8')
 
 
 # ============ 错误处理 ============
@@ -102,8 +89,15 @@ class DataTransformer:
     - fields/stats: 数据信息
     """
 
-    def __init__(self, data: Optional[List[Dict[str, Any]]] = None):
+    def __init__(
+        self,
+        data: Optional[List[Dict[str, Any]]] = None,
+        _source_path: Optional[str] = None,
+        _lineage_tracker: Optional[LineageTracker] = None,
+    ):
         self._data = data if data is not None else []
+        self._source_path = _source_path
+        self._lineage_tracker = _lineage_tracker
 
     @property
     def data(self) -> List[Dict[str, Any]]:
@@ -122,23 +116,38 @@ class DataTransformer:
     # ============ 加载/保存 ============
 
     @classmethod
-    def load(cls, filepath: str) -> 'DataTransformer':
+    def load(cls, filepath: str, track_lineage: bool = False) -> 'DataTransformer':
         """
         从文件加载数据。
 
         支持格式: jsonl, json, csv, parquet（自动检测）
+
+        Args:
+            filepath: 文件路径
+            track_lineage: 是否追踪血缘（默认 False）
         """
         data = load_data(filepath)
-        return cls(data)
+        tracker = LineageTracker(filepath) if track_lineage else None
+        return cls(data, _source_path=filepath, _lineage_tracker=tracker)
 
-    def save(self, filepath: str) -> None:
+    def save(self, filepath: str, lineage: bool = False) -> None:
         """
         保存数据到文件。
 
         支持格式: jsonl, json, csv, parquet（根据扩展名）
+
+        Args:
+            filepath: 文件路径
+            lineage: 是否保存血缘元数据（默认 False）
         """
         save_data(self._data, filepath)
 
+        # 保存血缘记录
+        if lineage and self._lineage_tracker:
+            lineage_path = self._lineage_tracker.save(filepath, len(self._data))
+            import sys
+            print(f"📜 血缘记录已保存: {lineage_path}", file=sys.stderr)
+
     # ============ 核心转换 ============
 
     def to(
@@ -230,7 +239,16 @@ class DataTransformer:
             >>> # 原始模式（大数据集推荐）
             >>> dt.transform(lambda x: {"q": x["q"]}, raw=True).save("output.jsonl")
         """
-        return DataTransformer(self.to(func, on_error=on_error, raw=raw))
+        input_count = len(self._data)
+        result = self.to(func, on_error=on_error, raw=raw)
+        output_count = len(result)
+
+        # 传递血缘追踪器并记录操作
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("transform", {"func": func}, input_count, output_count)
+
+        return DataTransformer(result, _lineage_tracker=tracker)
 
     # ============ 数据筛选 ============
 
@@ -281,7 +299,12 @@ class DataTransformer:
         if errors:
             _print_error_summary(errors, len(self._data))
 
-        return DataTransformer(filtered)
+        # 传递血缘追踪器并记录操作
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("filter", {"func": func}, len(self._data), len(filtered))
+
+        return DataTransformer(filtered, _lineage_tracker=tracker)
 
     def sample(self, n: int, seed: Optional[int] = None) -> 'DataTransformer':
         """
@@ -295,16 +318,30 @@ class DataTransformer:
         if seed is not None:
             random.seed(seed)
 
+        input_count = len(self._data)
         data = self._data[:] if n >= len(self._data) else random.sample(self._data, n)
-        return DataTransformer(data)
+
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("sample", {"n": n, "seed": seed}, input_count, len(data))
+
+        return DataTransformer(data, _lineage_tracker=tracker)
 
     def head(self, n: int = 10) -> 'DataTransformer':
         """取前 n 条"""
-        return DataTransformer(self._data[:n])
+        data = self._data[:n]
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("head", {"n": n}, len(self._data), len(data))
+        return DataTransformer(data, _lineage_tracker=tracker)
 
     def tail(self, n: int = 10) -> 'DataTransformer':
         """取后 n 条"""
-        return DataTransformer(self._data[-n:])
+        data = self._data[-n:]
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("tail", {"n": n}, len(self._data), len(data))
+        return DataTransformer(data, _lineage_tracker=tracker)
 
     def dedupe(
         self,
@@ -338,7 +375,11 @@ class DataTransformer:
                 seen.add(k)
                 result.append(item)
 
-        return DataTransformer(result)
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("dedupe", {"key": key}, len(self._data), len(result))
+
+        return DataTransformer(result, _lineage_tracker=tracker)
 
     def _get_dedupe_key(
         self,
@@ -442,7 +483,17 @@ class DataTransformer:
 
         # 按原顺序保留数据
         result = [self._data[i] for i in sorted(keep_indices)]
-        return DataTransformer(result)
+
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record(
+                "dedupe_similar",
+                {"key": key, "threshold": threshold, "num_perm": num_perm, "ngram": ngram},
+                len(self._data),
+                len(result),
+            )
+
+        return DataTransformer(result, _lineage_tracker=tracker)
 
     def _get_text_for_similarity(
         self,
@@ -581,7 +632,12 @@ class DataTransformer:
         if seed is not None:
             random.seed(seed)
         random.shuffle(data)
-        return DataTransformer(data)
+
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("shuffle", {"seed": seed}, len(self._data), len(data))
+
+        return DataTransformer(data, _lineage_tracker=tracker)
 
     def split(self, ratio: float = 0.8, seed: Optional[int] = None) -> tuple:
         """
@@ -596,7 +652,16 @@ class DataTransformer:
         """
         data = self.shuffle(seed).data
         split_idx = int(len(data) * ratio)
-        return DataTransformer(data[:split_idx]), DataTransformer(data[split_idx:])
+
+        # 分割后血缘追踪器各自独立
+        tracker = self._lineage_tracker
+        if tracker:
+            tracker.record("split", {"ratio": ratio, "seed": seed}, len(self._data), len(data))
+
+        return (
+            DataTransformer(data[:split_idx], _lineage_tracker=tracker),
+            DataTransformer(data[split_idx:], _lineage_tracker=tracker),
+        )
 
     # ============ 并行处理 ============
 

dtflow/lineage.py

@@ -0,0 +1,407 @@
+"""
+数据血缘模块
+
+记录数据处理的完整历史，支持数据溯源和版本对比。
+"""
+import hashlib
+import os
+
+import orjson
+import platform
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+# 血缘元数据版本
+LINEAGE_VERSION = "1.0"
+
+# 元数据文件后缀
+LINEAGE_SUFFIX = ".lineage.json"
+
+
+def _get_file_hash(filepath: str, sample_size: int = 10000) -> str:
+    """
+    计算文件内容哈希（采样方式，避免大文件性能问题）。
+
+    Args:
+        filepath: 文件路径
+        sample_size: 采样字节数
+
+    Returns:
+        文件哈希值（前16位）
+    """
+    hasher = hashlib.sha256()
+    file_size = os.path.getsize(filepath)
+
+    with open(filepath, "rb") as f:
+        # 读取文件头
+        hasher.update(f.read(sample_size))
+
+        # 如果文件较大，还要读取中间和尾部
+        if file_size > sample_size * 3:
+            f.seek(file_size // 2)
+            hasher.update(f.read(sample_size))
+            f.seek(-sample_size, 2)
+            hasher.update(f.read(sample_size))
+
+    return hasher.hexdigest()[:16]
+
+
+def _get_lineage_path(data_path: str) -> str:
+    """获取血缘元数据文件路径"""
+    return str(data_path) + LINEAGE_SUFFIX
+
+
+def _get_environment_info() -> Dict[str, str]:
+    """获取运行环境信息"""
+    return {
+        "python_version": platform.python_version(),
+        "platform": platform.system(),
+        "hostname": platform.node(),
+        "user": os.environ.get("USER", os.environ.get("USERNAME", "unknown")),
+    }
+
+
+class LineageRecord:
+    """血缘记录"""
+
+    def __init__(
+        self,
+        source: Optional[str] = None,
+        operations: Optional[List[Dict[str, Any]]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        self.version = LINEAGE_VERSION
+        self.created_at = datetime.now().isoformat()
+        self.source = source
+        self.operations = operations or []
+        self.metadata = metadata or {}
+        self.environment = _get_environment_info()
+
+    def add_operation(
+        self,
+        op_type: str,
+        params: Optional[Dict[str, Any]] = None,
+        input_count: Optional[int] = None,
+        output_count: Optional[int] = None,
+    ) -> "LineageRecord":
+        """添加操作记录"""
+        op = {
+            "type": op_type,
+            "timestamp": datetime.now().isoformat(),
+        }
+        if params:
+            op["params"] = params
+        if input_count is not None:
+            op["input_count"] = input_count
+        if output_count is not None:
+            op["output_count"] = output_count
+
+        self.operations.append(op)
+        return self
+
+    def to_dict(self) -> Dict[str, Any]:
+        """转换为字典"""
+        return {
+            "version": self.version,
+            "created_at": self.created_at,
+            "source": self.source,
+            "operations": self.operations,
+            "metadata": self.metadata,
+            "environment": self.environment,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "LineageRecord":
+        """从字典创建"""
+        record = cls(
+            source=data.get("source"),
+            operations=data.get("operations", []),
+            metadata=data.get("metadata", {}),
+        )
+        record.version = data.get("version", LINEAGE_VERSION)
+        record.created_at = data.get("created_at", datetime.now().isoformat())
+        record.environment = data.get("environment", {})
+        return record
+
+
+class LineageTracker:
+    """
+    血缘追踪器
+
+    用于记录数据处理的完整历史。
+    """
+
+    def __init__(self, source_path: Optional[str] = None):
+        """
+        初始化追踪器。
+
+        Args:
+            source_path: 源数据文件路径
+        """
+        self.source_path = source_path
+        self.source_lineage = None
+        self.operations: List[Dict[str, Any]] = []
+
+        # 如果源文件有血缘记录，加载它
+        if source_path:
+            self.source_lineage = load_lineage(source_path)
+
+    def record(
+        self,
+        op_type: str,
+        params: Optional[Dict[str, Any]] = None,
+        input_count: Optional[int] = None,
+        output_count: Optional[int] = None,
+    ) -> "LineageTracker":
+        """
+        记录一次操作。
+
+        Args:
+            op_type: 操作类型 (filter, transform, dedupe, sample, etc.)
+            params: 操作参数
+            input_count: 输入数据量
+            output_count: 输出数据量
+
+        Returns:
+            self，支持链式调用
+        """
+        op = {
+            "type": op_type,
+            "timestamp": datetime.now().isoformat(),
+        }
+        if params:
+            # 清理参数，移除不可序列化的内容
+            op["params"] = _sanitize_params(params)
+        if input_count is not None:
+            op["input_count"] = input_count
+        if output_count is not None:
+            op["output_count"] = output_count
+
+        self.operations.append(op)
+        return self
+
+    def build_record(self, output_path: str, output_count: int) -> LineageRecord:
+        """
+        构建最终的血缘记录。
+
+        Args:
+            output_path: 输出文件路径
+            output_count: 输出数据量
+
+        Returns:
+            LineageRecord 对象
+        """
+        # 构建来源信息
+        source_info = None
+        if self.source_path:
+            source_info = {
+                "path": str(self.source_path),
+                "hash": _get_file_hash(self.source_path) if os.path.exists(self.source_path) else None,
+            }
+            # 如果源文件有血缘，记录血缘链
+            if self.source_lineage:
+                source_info["lineage_ref"] = _get_lineage_path(self.source_path)
+
+        record = LineageRecord(
+            source=source_info,
+            operations=self.operations,
+            metadata={
+                "output_path": str(output_path),
+                "output_count": output_count,
+            },
+        )
+
+        return record
+
+    def save(self, output_path: str, output_count: int) -> str:
+        """
+        保存血缘记录到文件。
+
+        Args:
+            output_path: 输出数据文件路径
+            output_count: 输出数据量
+
+        Returns:
+            血缘文件路径
+        """
+        record = self.build_record(output_path, output_count)
+        lineage_path = _get_lineage_path(output_path)
+
+        with open(lineage_path, "wb") as f:
+            f.write(orjson.dumps(record.to_dict(), option=orjson.OPT_INDENT_2))
+
+        return lineage_path
+
+
+def _sanitize_params(params: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    清理参数，移除不可序列化的内容。
+    """
+    result = {}
+    for key, value in params.items():
+        if callable(value):
+            # 函数：只记录名称
+            result[key] = f"<function:{getattr(value, '__name__', 'anonymous')}>"
+        elif isinstance(value, (str, int, float, bool, type(None))):
+            result[key] = value
+        elif isinstance(value, (list, tuple)):
+            result[key] = [_sanitize_value(v) for v in value]
+        elif isinstance(value, dict):
+            result[key] = _sanitize_params(value)
+        else:
+            result[key] = str(value)
+    return result
+
+
+def _sanitize_value(value: Any) -> Any:
+    """清理单个值"""
+    if callable(value):
+        return f"<function:{getattr(value, '__name__', 'anonymous')}>"
+    elif isinstance(value, (str, int, float, bool, type(None))):
+        return value
+    elif isinstance(value, dict):
+        return _sanitize_params(value)
+    else:
+        return str(value)
+
+
+# ============ 公共 API ============
+
+
+def load_lineage(data_path: str) -> Optional[LineageRecord]:
+    """
+    加载数据文件的血缘记录。
+
+    Args:
+        data_path: 数据文件路径
+
+    Returns:
+        LineageRecord 或 None（如果没有血缘记录）
+    """
+    lineage_path = _get_lineage_path(data_path)
+    if not os.path.exists(lineage_path):
+        return None
+
+    try:
+        with open(lineage_path, "rb") as f:
+            data = orjson.loads(f.read())
+        return LineageRecord.from_dict(data)
+    except (orjson.JSONDecodeError, IOError):
+        return None
+
+
+def get_lineage_chain(data_path: str, max_depth: int = 10) -> List[LineageRecord]:
+    """
+    获取完整的血缘链。
+
+    Args:
+        data_path: 数据文件路径
+        max_depth: 最大追溯深度
+
+    Returns:
+        血缘记录列表，从最新到最旧
+    """
+    chain = []
+    current_path = data_path
+    visited = set()
+
+    for _ in range(max_depth):
+        if current_path in visited:
+            break  # 避免循环引用
+        visited.add(current_path)
+
+        record = load_lineage(current_path)
+        if not record:
+            break
+
+        chain.append(record)
+
+        # 追溯到源文件
+        if record.source and isinstance(record.source, dict):
+            source_path = record.source.get("path")
+            if source_path and os.path.exists(source_path):
+                current_path = source_path
+            else:
+                break
+        else:
+            break
+
+    return chain
+
+
+def format_lineage_report(data_path: str) -> str:
+    """
+    格式化血缘报告。
+
+    Args:
+        data_path: 数据文件路径
+
+    Returns:
+        格式化的报告字符串
+    """
+    chain = get_lineage_chain(data_path)
+
+    if not chain:
+        return f"文件 {data_path} 没有血缘记录"
+
+    lines = []
+    lines.append(f"📊 数据血缘报告: {data_path}")
+    lines.append("=" * 60)
+
+    for i, record in enumerate(chain):
+        prefix = "└─" if i == len(chain) - 1 else "├─"
+        indent = "  " * i
+
+        # 基本信息
+        lines.append(f"{indent}{prefix} 版本 {i + 1}")
+        lines.append(f"{indent}   创建时间: {record.created_at}")
+
+        # 来源信息
+        if record.source:
+            if isinstance(record.source, dict):
+                lines.append(f"{indent}   来源: {record.source.get('path', 'unknown')}")
+                if record.source.get("hash"):
+                    lines.append(f"{indent}   哈希: {record.source['hash']}")
+            else:
+                lines.append(f"{indent}   来源: {record.source}")
+
+        # 操作列表
+        if record.operations:
+            lines.append(f"{indent}   操作链:")
+            for j, op in enumerate(record.operations):
+                op_prefix = "└─" if j == len(record.operations) - 1 else "├─"
+                op_type = op.get("type", "unknown")
+                input_count = op.get("input_count", "?")
+                output_count = op.get("output_count", "?")
+                lines.append(f"{indent}     {op_prefix} {op_type}: {input_count} → {output_count}")
+
+                # 显示参数
+                if op.get("params"):
+                    for key, value in op["params"].items():
+                        lines.append(f"{indent}        {key}: {value}")
+
+        # 元数据
+        if record.metadata:
+            output_count = record.metadata.get("output_count")
+            if output_count:
+                lines.append(f"{indent}   输出数量: {output_count}")
+
+        lines.append("")
+
+    return "\n".join(lines)
+
+
+def has_lineage(data_path: str) -> bool:
+    """检查文件是否有血缘记录"""
+    return os.path.exists(_get_lineage_path(data_path))
+
+
+def delete_lineage(data_path: str) -> bool:
+    """删除血缘记录"""
+    lineage_path = _get_lineage_path(data_path)
+    if os.path.exists(lineage_path):
+        os.remove(lineage_path)
+        return True
+    return False

dtflow/mcp/cli.py

@@ -3,7 +3,7 @@
 提供 MCP 服务的安装和管理命令。
 """
 
-import json
+import orjson
 import os
 import platform
 from pathlib import Path
@@ -63,9 +63,9 @@ def _install_to_config(config_path: Path, name: str, target_name: str) -> bool:
     config = {}
     if config_path.exists():
         try:
-            with open(config_path, "r", encoding="utf-8") as f:
-                config = json.load(f)
-        except json.JSONDecodeError:
+            with open(config_path, "rb") as f:
+                config = orjson.loads(f.read())
+        except orjson.JSONDecodeError:
             if console:
                 console.print(f"[yellow]警告:[/yellow] {target_name} 配置文件格式错误，将创建新配置")
             else:
@@ -87,8 +87,8 @@ def _install_to_config(config_path: Path, name: str, target_name: str) -> bool:
 
     # 写入配置
     try:
-        with open(config_path, "w", encoding="utf-8") as f:
-            json.dump(config, f, indent=2, ensure_ascii=False)
+        with open(config_path, "wb") as f:
+            f.write(orjson.dumps(config, option=orjson.OPT_INDENT_2))
         return True
     except Exception as e:
         if console:
@@ -108,9 +108,9 @@ def _uninstall_from_config(config_path: Path, name: str, target_name: str) -> bo
         return False
 
     try:
-        with open(config_path, "r", encoding="utf-8") as f:
-            config = json.load(f)
-    except json.JSONDecodeError:
+        with open(config_path, "rb") as f:
+            config = orjson.loads(f.read())
+    except orjson.JSONDecodeError:
         return False
 
     if "mcpServers" not in config or name not in config["mcpServers"]:
@@ -119,8 +119,8 @@ def _uninstall_from_config(config_path: Path, name: str, target_name: str) -> bo
     del config["mcpServers"][name]
 
     try:
-        with open(config_path, "w", encoding="utf-8") as f:
-            json.dump(config, f, indent=2, ensure_ascii=False)
+        with open(config_path, "wb") as f:
+            f.write(orjson.dumps(config, option=orjson.OPT_INDENT_2))
         return True
     except Exception:
         return False
@@ -141,9 +141,9 @@ def _show_config_status(config_path: Path, target_name: str):
         return
 
     try:
-        with open(config_path, "r", encoding="utf-8") as f:
-            config = json.load(f)
-    except json.JSONDecodeError:
+        with open(config_path, "rb") as f:
+            config = orjson.loads(f.read())
+    except orjson.JSONDecodeError:
         if console:
             console.print("  [red]配置文件格式错误[/red]")
         else: