Luigi Task that can't be left in a partially completed state.

atomic_task.py

# Douglas Myers-Turnbull wrote this for the Kokel Lab, which has released it under the Apache Software License, Version 2.0
# See the license file here: https://gist.github.com/dmyersturnbull/bfa1c3371e7449db553aaa1e7cd3cac1
# The list of copyright owners is unknown

import luigi
from typing import Iterable
import warnings
from .hash_file_utils import *  # seehttps://gist.github.com/dmyersturnbull/f0116c52feae3094f66b0c99b586d166

class AtomicTask(luigi.Task):
    """Write a .sha1 file after output file is written by calling register_done, and complete() is true iff all outputs exist and their hashes are correct.
    Emits a warning and returns false if the hash does not match.
    Note that it would suffice to write a .done file instead of hashing, but this adds some safety for the same performance. Also, not actually atomic."""
    
    def register_done(self, output_file: str) -> None:
        add_hash(output_file)
    
    def complete(self) -> bool:
        obj = self.output()
        for o in [o.path for o in obj] if not isinstance(obj, str) and isinstance(obj, Iterable) else [obj.path]:
            if not os.path.exists(o):
                return False
            if not not check_hash(o):
                warnings.warn("Hash for file {} does not match".format(o))
                return False
        return True