RenJW/RNAediting

1import os2import shutil3import subprocess4import hashlib5import glob6import re7from pathlib import Path8 9class DataHandler:10    """处理数据上传、下载、校验等操作"""11    12    def download_from_oss(self, access_key_id, access_key_secret, oss_path, target_dir):13        """14        从OSS下载数据15        16        Args:17            access_key_id: OSS访问密钥ID18            access_key_secret: OSS访问密钥19            oss_path: OSS路径20            target_dir: 目标目录21            22        Returns:23            str: 下载结果消息24        """25        try:26            # 创建 ossutil 配置文件27            config_content = f"""[Credentials]28language=CH29endpoint=oss-cn-hangzhou.aliyuncs.com30accessKeyID={access_key_id}31accessKeySecret={access_key_secret}32"""33            config_path = Path.home() / ".ossutilconfig"34            with open(config_path, 'w') as f:35                f.write(config_content)36            37            # 检查 ossutil64 是否存在38            ossutil_path = "./ossutil64"39            if not os.path.exists(ossutil_path):40                # 尝试使用系统路径41                ossutil_path = "ossutil64"42            else:43                # 设置执行权限44                import stat45                try:46                    os.chmod(ossutil_path, stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH | stat.S_IXOTH)47                except:48                    pass49            50            # 执行下载命令51            cmd = [ossutil_path, "cp", "-r", "--include", "*", oss_path, target_dir]52            53            result = subprocess.run(54                cmd,55                capture_output=True,56                text=True,57                timeout=3600  # 1小时超时58            )59            60            if result.returncode == 0:61                return f"✅ OSS数据下载成功\n{result.stdout}"62            else:63                return f"❌ OSS数据下载失败\n错误: {result.stderr}"64                65        except subprocess.TimeoutExpired:66            return "❌ 下载超时(超过1小时)"67        except Exception as e:68            return f"❌ 下载过程出错: {str(e)}"69    70    def check_md5(self, data_dir):71        """72        检查MD5校验和73        74        Args:75            data_dir: 数据目录76            77        Returns:78            str: 校验结果79        """80        try:81            # 查找 md5.md5 文件82            md5_files = list(Path(data_dir).rglob("md5.md5"))83            84            if not md5_files:85                return "⚠️ 未找到 md5.md5 文件,跳过MD5校验"86            87            md5_file = md5_files[0]88            89            # 读取MD5文件90            with open(md5_file, 'r') as f:91                md5_records = {}92                for line in f:93                    line = line.strip()94                    if line:95                        parts = line.split()96                        if len(parts) >= 2:97                            md5_hash = parts[0]98                            filename = parts[1]99                            md5_records[filename] = md5_hash100            101            # 校验每个文件102            md5_dir = md5_file.parent103            errors = []104            105            for filename, expected_md5 in md5_records.items():106                file_path = md5_dir / filename107                if file_path.exists():108                    actual_md5 = self._calculate_md5(file_path)109                    if actual_md5 != expected_md5:110                        errors.append(f"  ❌ {filename}: MD5不匹配")111                else:112                    errors.append(f"  ❌ {filename}: 文件不存在")113            114            if errors:115                raise Exception("MD5有异常\n" + "\n".join(errors))116            117            return f"✅ MD5校验通过 (校验了 {len(md5_records)} 个文件)"118            119        except Exception as e:120            if "MD5有异常" in str(e):121                raise122            return f"⚠️ MD5校验过程出错: {str(e)}"123    124    def _calculate_md5(self, file_path):125        """计算文件的MD5值"""126        md5_hash = hashlib.md5()127        with open(file_path, 'rb') as f:128            for chunk in iter(lambda: f.read(4096), b""):129                md5_hash.update(chunk)130        return md5_hash.hexdigest()131    132    def move_fastq_files(self, data_dir):133        """134        将嵌套目录中的 fastq 文件移动到 00-data 根目录135        136        Args:137            data_dir: 数据目录 (00-data)138            139        Returns:140            str: 移动结果141        """142        try:143            data_path = Path(data_dir)144            145            # 查找所有 .fq.gz 和 .fastq.gz 文件146            fastq_files = []147            fastq_files.extend(data_path.rglob("*.fq.gz"))148            fastq_files.extend(data_path.rglob("*.fastq.gz"))149            150            if not fastq_files:151                return f"⚠️ 未找到任何测序文件 (.fq.gz 或 .fastq.gz)\n   请检查下载是否成功"152            153            moved_count = 0154            for file_path in fastq_files:155                # 如果文件不在根目录,则移动156                if file_path.parent != data_path:157                    target_path = data_path / file_path.name158                    shutil.move(str(file_path), str(target_path))159                    moved_count += 1160            161            # 清理空目录162            for item in data_path.iterdir():163                if item.is_dir():164                    try:165                        # 删除空目录及其子目录166                        if not any(item.rglob("*")):167                            shutil.rmtree(item)168                    except:169                        pass170            171            return f"✅ 移动了 {moved_count} 个测序文件到根目录"172            173        except Exception as e:174            return f"⚠️ 移动文件时出错: {str(e)}"175    176    def extract_sample_names(self, data_dir):177        """178        从文件名中提取样本名179        180        Args:181            data_dir: 数据目录182            183        Returns:184            list: 样本名列表185        """186        data_path = Path(data_dir)187        samples = set()188        189        # 查找所有测序文件190        for file_path in data_path.glob("*.fq.gz"):191            filename = file_path.name192            193            # 匹配 _1.fq.gz 或 _2.fq.gz 模式194            match = re.match(r"(.+)_[12]\.fq\.gz$", filename)195            if match:196                sample_name = match.group(1)197                samples.add(sample_name)198        199        # 也支持 .fastq.gz 后缀200        for file_path in data_path.glob("*.fastq.gz"):201            filename = file_path.name202            match = re.match(r"(.+)_[12]\.fastq\.gz$", filename)203            if match:204                sample_name = match.group(1)205                samples.add(sample_name)206        207        return sorted(list(samples))208    209    def validate_paired_files(self, data_dir, samples):210        """211        验证每个样本都有配对的测序文件212        213        Args:214            data_dir: 数据目录215            samples: 样本名列表216            217        Returns:218            tuple: (是否通过, 错误消息)219        """220        data_path = Path(data_dir)221        errors = []222        223        for sample in samples:224            file1 = data_path / f"{sample}_1.fq.gz"225            file2 = data_path / f"{sample}_2.fq.gz"226            227            if not file1.exists():228                errors.append(f"缺少文件: {file1.name}")229            if not file2.exists():230                errors.append(f"缺少文件: {file2.name}")231        232        if errors:233            return False, "\n".join(errors)234        235        return True, "所有样本文件配对完整"236