diff --git a/src/utils/util.py b/src/utils/util.py index ce4a989..0d9fea2 100644 --- a/src/utils/util.py +++ b/src/utils/util.py @@ -32,6 +32,70 @@ def timeit(f): # decorator return wraper +def is_file_complete(file_path) -> bool: + """检查文件是否已完整生成(需配套同名 .done 侧标使用)。""" + from pathlib import Path + target = Path(file_path) + done_marker = target.with_suffix(target.suffix + '.done') + return target.exists() and done_marker.exists() + + +def mark_file_complete(file_path) -> None: + """为成功生成的文件打上完整性标记。 + + 原子化约定:仅在步骤真正无异常执行完、产物确认落盘后再调用; + “存在即跳过”的判定应改为 is_file_complete(),杜绝半成品复用。 + """ + from pathlib import Path + target = Path(file_path) + done_marker = target.with_suffix(target.suffix + '.done') + done_marker.touch(exist_ok=True) + + +def atomic_filepath(final_path: str, tmp_suffix: str = '.__wip'): + """原子写入上下文管理器:文件要么完整生成,要么完全不出现。 + + 工作流: + 1. 生成带 .__wip (Work In Progress) 后缀的临时路径供内部写入。 + 2. 中途异常 → 自动删除 .__wip,绝不污染目标路径。 + 3. 仅代码块无异常执行完毕 → 同卷原子重命名 + 自动打 .done 标记。 + + 用法: + with atomic_filepath(out_path) as tmp: + result_df.to_csv(tmp, index=False) + """ + from contextlib import contextmanager + from pathlib import Path + + @contextmanager + def _inner(): + final_p = Path(final_path) + tmp_p = final_p.with_suffix(final_p.suffix + tmp_suffix) + + # 清除历史遗留的死锁临时文件 + if tmp_p.exists(): + try: + tmp_p.unlink() + except OSError: + pass + try: + yield str(tmp_p) + # --- 无异常走完,执行提交 (Commit) --- + if tmp_p.exists(): + os.replace(tmp_p, final_p) # 同卷原子替换 + mark_file_complete(str(final_p)) # 自动打标 + except Exception as e: + # --- 发生异常,执行回滚 (Rollback) --- + if tmp_p.exists(): + try: + tmp_p.unlink() + except OSError: + pass + raise e + + return _inner() + + def get_hdr_file_path(file_path): """获取 ENVI 头文件路径(鲁棒版:多命名规范兼容) @@ -247,6 +311,29 @@ def write_bands(imgpath_in, imgpath_out, *args): del dataset, dst_ds +def create_output_dataset(path, width, height, n_bands, dtype): + """创建(可能超大的)多波段输出数据集。 + + ENVI driver 对"每个波段字节数 > 2GB 且波段数 > 1"的 BSQ 输出存在 + int32 偏移溢出(GDAL 内部用 int 存第 2 个波段起的数据偏移,会抛 + "Int overflow occurred")。此类大图自动改用 GTiff + BIGTIFF=YES; + 常规尺寸仍用 ENVI(.bsq 惯例),保持与历史一致。 + + 返回: 已创建且可写的 gdal.Dataset;失败抛 RuntimeError + """ + band_bytes = max(1, gdal.GetDataTypeSize(dtype) // 8) + per_band = width * height * band_bytes + if n_bands > 1 and per_band > 2_000_000_000: + ds = gdal.GetDriverByName('GTiff').Create( + path, width, height, n_bands, dtype, + options=['BIGTIFF=YES']) + else: + ds = gdal.GetDriverByName('ENVI').Create(path, width, height, n_bands, dtype) + if ds is None: + raise RuntimeError(f"无法创建输出文件: {path}") + return ds + + def append2filename(file_path, txt2add): imgfile_out_tmp = os.path.splitext(file_path) new_file_path = imgfile_out_tmp[0] + "_" + txt2add + imgfile_out_tmp[1]